megamoe

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Generalize
2026-06-22 23:59:53 +00:00 · 2026-06-22 19:19:36 +00:00 · 2026-06-22 04:57:51 +00:00 · 2026-06-22 04:52:12 +00:00 · 2026-06-22 04:51:16 +00:00 · 2026-06-22 01:51:39 +00:00
203 changed files with 6952 additions and 1449 deletions
@@ -647,7 +647,7 @@ steps:
  - pytest -v -s v1/cudagraph/test_cudagraph_mode.py

 - label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 35
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
  agent_pool: mi250_1
  optional: true
@@ -1594,9 +1594,10 @@ steps:
 #----------------------------------------------------------  mi300 · kernels  ----------------------------------------------------------#

 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 55
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  parallelism: 2
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
@@ -1627,10 +1628,11 @@ steps:
  - pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py  kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py

 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
-  parallelism: 4
+  optional: true
+  parallelism: 5
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -2075,19 +2077,6 @@ steps:
  - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
  - pytest -v -s v1/spec_decode/test_acceptance_length.py -m slow_test

- label: e2e Core (1 GPU) # TBD
-  timeout_in_minutes: 180
-  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
-  agent_pool: mi300_1
-  optional: true
-  working_dir: "/vllm-workspace/tests"
-  source_file_dependencies:
-  - vllm/v1/
-  - tests/v1/e2e/
-  - vllm/platforms/rocm.py
-  commands:
-  - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
-
 - label: e2e Scheduling (1 GPU) # TBD
  timeout_in_minutes: 180
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
@@ -2133,9 +2122,10 @@ steps:
  - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"

 - label: Spec Decode Eagle # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 45
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
  agent_pool: mi300_1
+  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - vllm/v1/spec_decode/
@@ -3053,7 +3043,7 @@ steps:
 #----------------------------------------------------------  mi355 · kernels  ----------------------------------------------------------#

 - label: Kernels (B200-MI355) # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 15
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  working_dir: "/vllm-workspace/"
@@ -3077,11 +3067,10 @@ steps:
  - pytest -v -s tests/kernels/attention/test_attention_selector.py

 - label: Kernels Attention Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 60
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
  parallelism: 2
-  optional: true
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - csrc/attention/
@@ -3095,10 +3084,10 @@ steps:
  - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

 - label: Kernels MoE Test %N # TBD
-  timeout_in_minutes: 180
+  timeout_in_minutes: 50
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
  agent_pool: mi355_1
-  parallelism: 4
+  parallelism: 5
  working_dir: "/vllm-workspace/tests"
  source_file_dependencies:
  - csrc/quantization/cutlass_w8a8/moe/
@@ -74,6 +74,16 @@ steps:
    - tests/v1/e2e/general/
  commands:
    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
+  mirror:
+    amd:
+      device: mi250_1
+      timeout_in_minutes: 35
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/
+      - tests/v1/e2e/general/
+      - vllm/platforms/rocm.py

 - label: V1 e2e (2 GPUs)
  key: v1-e2e-2-gpus
@@ -74,6 +74,20 @@ steps:
  commands:
    - pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
  parallelism: 2
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 55
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/attention/
+      - vllm/v1/attention
+      - vllm/model_executor/layers/attention
+      - tests/kernels/attention
+      - vllm/_aiter_ops.py
+      - vllm/envs.py
+      - vllm/platforms/rocm.py

 - label: Kernels Attention DiffKV Test (H100)
  key: kernels-attention-diffkv-test-h100
@@ -104,6 +118,7 @@ steps:
      source_file_dependencies:
      - csrc/quantization/
      - vllm/model_executor/layers/quantization
+      - vllm/config/
      - tests/kernels/quantization
      - tests/kernels/quantization/test_rocm_skinny_gemms.py
      - vllm/_aiter_ops.py
@@ -127,6 +142,22 @@ steps:
    - 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
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 50
+      source_file_dependencies:
+      - csrc/quantization/cutlass_w8a8/moe/
+      - csrc/moe/
+      - tests/kernels/moe
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/distributed/device_communicators/
+      - vllm/envs.py
+      - vllm/config
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      depends_on:
+      - image-build-amd

 - label: Kernels Mamba Test
  key: kernels-mamba-test
@@ -101,6 +101,16 @@ steps:
  num_devices: 8
  commands:
    - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
+  mirror:
+    amd:
+      device: mi300_8
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd
+      commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time
+      - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt

 - label: MoE Refactor Integration Test (H100 - TEMPORARY)
  key: moe-refactor-integration-test-h100-temporary
@@ -68,7 +68,6 @@ steps:
    - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model  # Otherwise, mp_method="spawn" doesn't work
  mirror:
    amd:
-      soft_fail: true
      device: mi325_1
      depends_on:
      - image-build-amd
@@ -107,6 +107,12 @@ steps:
    - tests/compile/passes
  commands:
  - pytest -s -v compile/passes --ignore compile/passes/distributed
+  mirror:
+    amd:
+      device: mi300_1
+      timeout_in_minutes: 180
+      depends_on:
+      - image-build-amd

 - label: PyTorch Fullgraph Smoke Test
  key: pytorch-fullgraph-smoke-test
@@ -12,6 +12,20 @@ steps:
    - tests/v1/e2e/spec_decode/
  commands:
    - pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
+  mirror:
+    amd:
+      device: mi325_1
+      timeout_in_minutes: 45
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/v1/spec_decode/
+      - vllm/v1/worker/gpu/spec_decode/
+      - vllm/model_executor/model_loader/
+      - vllm/v1/sample/
+      - vllm/model_executor/layers/
+      - tests/v1/e2e/spec_decode/
+      - vllm/platforms/rocm.py

 - label: Spec Decode Eagle Nightly B200
  key: spec-decode-eagle-nightly-b200
@@ -0,0 +1,35 @@
+---
+name: ci-fails-buildkite
+description: Fetch and diagnose vLLM Buildkite CI failure logs. Use when investigating failing CI jobs on a PR or build, when the user pastes a buildkite.com URL, or asks to fetch/diagnose CI logs.
+---
+
+# Diagnosing vLLM Buildkite CI Failures
+
+Buildkite logs are public; no login needed.
+
+`.buildkite/scripts/ci-fetch-log.sh` saves each log as `ci-<build>-<job-name>.log`, stripped of timestamps and ANSI codes. Existing files are kept; set `CI_FETCH_LOG_FORCE=1` to refetch.
+
+## Fetching logs
+
+```bash
+# All failed jobs in a PR's latest build (current branch's PR if omitted):
+.buildkite/scripts/ci-fetch-log.sh --pr <PR>
+
+# All failed jobs in a build (--soft also includes soft-failed jobs;
+# --all fetches every finished job):
+.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>"
+
+# One job — `gh pr checks` URLs (#<job_uuid>) and web UI URLs (?sid=) both
+# work; pass "-" as a second argument to stream to stdout:
+.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
+```
+
+To clean an already-downloaded log with `.buildkite/scripts/ci-clean-log.sh`:
+
+```bash
+./ci-clean-log.sh ci.log
+```
+
+## Reference
+
+See [docs/contributing/ci/failures.md](../../../docs/contributing/ci/failures.md) for the full guide: filing CI failure issues, investigating/bisecting, reproducing flaky tests, and daily triage.
@@ -2,15 +2,14 @@
 # for more info about CODEOWNERS file

 # This lists cover the "core" components of vLLM that require careful review
-/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng @vadiklyutiy
+/vllm/compilation @zou3519 @youkaichao @ProExpertProg @BoyuanFeng
 /vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery @xuechendi
 /vllm/lora @jeejeelee
 /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni
 /vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye
 /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye
 /vllm/model_executor/layers/mamba @tdoublep @tomeras91
-/vllm/model_executor/layers/mamba/gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
-/vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
+/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy
 /vllm/model_executor/model_loader @22quinn
 /vllm/model_executor/layers/batch_invariant.py @yewentao256 
 /vllm/ir @ProExpertProg
@@ -199,7 +199,9 @@ cython_debug/
 .vscode/

 # Claude
-.claude/
+.claude/*
+!.claude/skills/
+!.claude/skills/**

 # Codex
 .codex/
@@ -114,17 +114,6 @@ Follow these rules for all code changes in this repository:
 - Keep comments and docstrings minimal and concise.
 - Assume the reader is familiar with vLLM.

-### Diagnosing CI failures
-
-Buildkite logs are public; no login needed. Details: [docs/contributing/ci/failures.md](docs/contributing/ci/failures.md).
-
-```bash
-# All failed-job logs for a PR's latest build (current branch's PR if omitted):
-.buildkite/scripts/ci-fetch-log.sh --pr <PR>
-# Any Buildkite build or job URL also works:
-.buildkite/scripts/ci-fetch-log.sh "<buildkite_url>"
-```
-
 ### Commit messages

 Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
@@ -319,82 +319,35 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
 endif()

 #
-# _C extension
+# Legacy _C extension (ROCm only — CUDA ops migrated to _C_stable_libtorch)
 #

-set(VLLM_EXT_SRC
-  "csrc/quantization/activation_kernels.cu"
-  "csrc/torch_bindings.cpp")
-
-if(VLLM_GPU_LANG STREQUAL "CUDA")
-  SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
-
-  # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
-  set(CUTLASS_REVISION "v4.4.2")
-
-  # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
-  if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
-    set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
-  endif()
-
-  if(VLLM_CUTLASS_SRC_DIR)
-    if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
-      get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
-    endif()
-    message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
-    FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
-  else()
-    FetchContent_Declare(
-        cutlass
-        GIT_REPOSITORY https://github.com/nvidia/cutlass.git
-        # Please keep this in sync with CUTLASS_REVISION line above.
-        GIT_TAG ${CUTLASS_REVISION}
-        GIT_PROGRESS TRUE
-
-        # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
-        # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
-        # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
-        GIT_SHALLOW TRUE
-    )
-  endif()
-  FetchContent_MakeAvailable(cutlass)
-
-  set_gencode_flags_for_srcs(
-    SRCS "${VLLM_EXT_SRC}"
-    CUDA_ARCHS "${CUDA_ARCHS}")
-
-# if CUDA endif
-endif()
-
-if (VLLM_GPU_LANG STREQUAL "HIP")
-  # Add QuickReduce kernels (ROCm-only; not part of stable ABI migration).
-  # TODO: Remove the cuda_view when ROCm upgrade to torch 2.11.
-  list(APPEND VLLM_EXT_SRC
+if(VLLM_GPU_LANG STREQUAL "HIP")
+  set(VLLM_EXT_SRC
+    "csrc/torch_bindings.cpp"
    "csrc/custom_quickreduce.cu"
    "csrc/cuda_view.cu"
-    "csrc/libtorch_stable/cuda_utils_kernels.cu"
-  )
-# if ROCM endif
-endif()
+    "csrc/libtorch_stable/cuda_utils_kernels.cu")

-message(STATUS "Enabling C extension.")
-define_extension_target(
-  _C
-  DESTINATION vllm
-  LANGUAGE ${VLLM_GPU_LANG}
-  SOURCES ${VLLM_EXT_SRC}
-  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
-  ARCHITECTURES ${VLLM_GPU_ARCHES}
-  INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
-  INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
-  USE_SABI 3
-  WITH_SOABI)
+  message(STATUS "Enabling C extension.")
+  define_extension_target(
+    _C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${VLLM_EXT_SRC}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
+    INCLUDE_DIRECTORIES ${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
+    USE_SABI 3
+    WITH_SOABI)

-# If CUTLASS is compiled on NVCC >= 12.5, it by default uses
-# cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
-# driver API. This causes problems when linking with earlier versions of CUDA.
-# Setting this variable sidesteps the issue by calling the driver directly.
-target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+  # If CUTLASS is compiled on NVCC >= 12.5, it by default uses
+  # cudaGetDriverEntryPointByVersion as a wrapper to avoid directly calling the
+  # driver API. This causes problems when linking with earlier versions of CUDA.
+  # Setting this variable sidesteps the issue by calling the driver directly.
+  target_compile_definitions(_C PRIVATE CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+endif() # _C HIP endif

 if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  #
@@ -403,6 +356,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  set(VLLM_STABLE_EXT_SRC
    "csrc/libtorch_stable/torch_bindings.cpp"
    "csrc/libtorch_stable/activation_kernels.cu"
+    "csrc/libtorch_stable/quantization/activation_kernels.cu"
    "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"
    "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"
    "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
@@ -429,6 +383,38 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

  if(VLLM_GPU_LANG STREQUAL "CUDA")
+    SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
+
+    # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.
+    set(CUTLASS_REVISION "v4.4.2")
+
+    # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
+    if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
+      set(VLLM_CUTLASS_SRC_DIR $ENV{VLLM_CUTLASS_SRC_DIR})
+    endif()
+
+    if(VLLM_CUTLASS_SRC_DIR)
+      if(NOT IS_ABSOLUTE VLLM_CUTLASS_SRC_DIR)
+        get_filename_component(VLLM_CUTLASS_SRC_DIR "${VLLM_CUTLASS_SRC_DIR}" ABSOLUTE)
+      endif()
+      message(STATUS "The VLLM_CUTLASS_SRC_DIR is set, using ${VLLM_CUTLASS_SRC_DIR} for compilation")
+      FetchContent_Declare(cutlass SOURCE_DIR ${VLLM_CUTLASS_SRC_DIR})
+    else()
+      FetchContent_Declare(
+          cutlass
+          GIT_REPOSITORY https://github.com/nvidia/cutlass.git
+          # Please keep this in sync with CUTLASS_REVISION line above.
+          GIT_TAG ${CUTLASS_REVISION}
+          GIT_PROGRESS TRUE
+
+          # Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
+          # Important: If GIT_SHALLOW is enabled then GIT_TAG works only with branch names and tags.
+          # So if the GIT_TAG above is updated to a commit hash, GIT_SHALLOW must be set to FALSE
+          GIT_SHALLOW TRUE
+      )
+    endif()
+    FetchContent_MakeAvailable(cutlass)
+
    list(APPEND VLLM_STABLE_EXT_SRC
      "csrc/libtorch_stable/cuda_view.cu"
      "csrc/libtorch_stable/cuda_utils_kernels.cu"
@@ -929,7 +915,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${FP4_SM120_SRCS}"
      CUDA_ARCHS "${FP4_SM120_ARCHS}")
    list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
    message(STATUS "Building SM12x NVFP4 for archs: ${FP4_SM120_ARCHS}")
@@ -962,7 +947,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${FP4_SM100_SRCS}"
      CUDA_ARCHS "${FP4_SM100_ARCHS}")
    list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
-    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
    message(STATUS "Building SM10x/11x NVFP4/MXFP4 for archs: ${FP4_SM100_ARCHS}")
@@ -60,6 +60,7 @@ endif()

 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
  set(QUTLASS_SOURCES
+    csrc/qutlass_registration.cpp
    ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
    ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
    ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
@@ -78,8 +79,19 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)

  if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
    list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
+    if(CUTLASS_TOOLS_UTIL_INCLUDE_DIR AND
+       EXISTS "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}/cutlass/util/packed_stride.hpp")
+      list(APPEND QUTLASS_INCLUDES "${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}")
+    else()
+      get_filename_component(_qutlass_cutlass_root "${CUTLASS_INCLUDE_DIR}" DIRECTORY)
+      if(EXISTS "${_qutlass_cutlass_root}/tools/util/include/cutlass/util/packed_stride.hpp")
+        list(APPEND QUTLASS_INCLUDES "${_qutlass_cutlass_root}/tools/util/include")
+      endif()
+    endif()
  elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
-    list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
+    list(APPEND QUTLASS_INCLUDES
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include"
+      "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/tools/util/include")
    message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
  else()
    message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
@@ -91,12 +103,23 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    CUDA_ARCHS "${QUTLASS_ARCHS}"
  )

-  target_sources(_C PRIVATE ${QUTLASS_SOURCES})
-  target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
-  target_compile_definitions(_C PRIVATE
+  # QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
+  # Keep it as its own extension (registers torch.ops._qutlass_C).
+  define_extension_target(
+    _qutlass_C
+    DESTINATION vllm
+    LANGUAGE ${VLLM_GPU_LANG}
+    SOURCES ${QUTLASS_SOURCES}
+    COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+    ARCHITECTURES ${VLLM_GPU_ARCHES}
+    INCLUDE_DIRECTORIES ${QUTLASS_INCLUDES}
+    USE_SABI 3
+    WITH_SOABI)
+
+  target_compile_definitions(_qutlass_C PRIVATE
    QUTLASS_DISABLE_PYBIND=1
    TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-  )
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

  set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
    $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -111,4 +134,5 @@ else()
      "[QUTLASS] Skipping build: no supported arch (12.0f / 10.0f) found in "
      "CUDA_ARCHS='${CUDA_ARCHS}'.")
  endif()
+  add_custom_target(_qutlass_C)
 endif()
@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
  using TileShapeD = typename MlaSm100Type::TileShapeD;
  arguments.problem_shape =
      cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
-  // Assumes device 0 when getting sm_count.
-  arguments.hw_info.sm_count =
-      sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
+  if (sm_count <= 0) {
+    int current_device = 0;
+    cudaGetDevice(&current_device);
+    arguments.hw_info.sm_count =
+        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
+  } else {
+    arguments.hw_info.sm_count = sm_count;
+  }
  arguments.split_kv = static_cast<int>(num_kv_splits);
  MlaSm100Type::Fmha::set_split_kv(arguments);

@@ -9,7 +9,7 @@
 #include <torch/headeronly/core/ScalarType.h>

 #include "../../cuda_compat.h"
-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "libtorch_stable/dispatch_utils.h"
 #include "libtorch_stable/torch_utils.h"

@@ -2,9 +2,25 @@

 #include <torch/csrc/stable/library.h>
 #include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/util/Exception.h>

 #include <optional>
 #include <string>
+#include <vector>
+
+#include <torch/csrc/stable/ops.h>
+
+inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
+  // Ensure tensor is on CUDA
+  STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
+
+  // Get the raw data pointer
+  void* data_ptr = tensor.mutable_data_ptr();
+
+  /// Create a new tensor from the raw data pointer
+  return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
+                                  tensor.device(), tensor.scalar_type());
+}

 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
                               torch::stable::Tensor& output_q,
@@ -371,6 +387,18 @@ void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void silu_and_mul_clamp(torch::stable::Tensor& out,
                        torch::stable::Tensor& input, double limit,
                        double alpha = 1.0, double beta = 0.0);
+
+void silu_and_mul_quant(torch::stable::Tensor& out,
+                        torch::stable::Tensor& input,
+                        torch::stable::Tensor& scale);
+
+void persistent_masked_m_silu_mul_quant(
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
+    bool use_ue8m0);
+
 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
 void gelu_tanh_and_mul(torch::stable::Tensor& out,
@@ -1,16 +1,12 @@
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/all.h>
-#include <c10/cuda/CUDAGuard.h>
+#include "libtorch_stable/torch_utils.h"

 #include <cmath>
-#include "core/math.hpp"
-#include "../cuda_compat.h"
-#include "dispatch_utils.h"

+#include "libtorch_stable/core/math.hpp"
+#include "cuda_compat.h"
+#include "libtorch_stable/dispatch_utils.h"
 #include "quantization/w8a8/fp8/common.cuh"

-#include <c10/util/Float8_e4m3fn.h>
-
 #ifndef USE_ROCM
  #include <cuda_bf16.h>
  #include <cuda_fp16.h>
@@ -33,7 +29,6 @@ typedef __hip_fp8x4_e4m3_fnuz __nv_fp8x4_e4m3;
  #endif
 #endif

-#include "core/registration.h"
 namespace vllm {

 template <typename T>
@@ -564,41 +559,47 @@ __global__ void silu_mul_fp8_quant_deep_gemm_kernel(
 }  // namespace vllm

 // Launch activation, gating, and quantize kernel.
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                               \
-  int d = input.size(-1) / 2;                                               \
-  int64_t num_tokens = input.numel() / input.size(-1);                      \
-  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);     \
-  dim3 block(std::min(d, 512));                                             \
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));         \
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();             \
-  VLLM_DISPATCH_FLOATING_TYPES(                                             \
-      input.scalar_type(), "act_and_mul_kernel", [&] {                      \
-        VLLM_DISPATCH_FP8_TYPES(                                            \
-            out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] {  \
-              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,    \
-                                             fp8_t>                         \
-                  <<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),       \
-                                               input.data_ptr<scalar_t>(),  \
-                                               scale.data_ptr<float>(), d); \
-            });                                                             \
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL)                             \
+  int d = input.size(-1) / 2;                                             \
+  int64_t num_tokens = input.numel() / input.size(-1);                    \
+  dim3 grid(num_tokens, num_tokens > 16 ? num_tokens > 32 ? 1 : 2 : 4);   \
+  dim3 block(std::min(d, 512));                                           \
+  const torch::stable::accelerator::DeviceGuard device_guard(             \
+      input.get_device_index());                                          \
+  const cudaStream_t stream =                                             \
+      get_current_cuda_stream(input.get_device_index());                  \
+  VLLM_STABLE_DISPATCH_FLOATING_TYPES(                                    \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                    \
+        VLLM_STABLE_DISPATCH_FP8_TYPES(                                   \
+            out.scalar_type(), "act_and_mul_quant_kernel_fp8_type", [&] { \
+              vllm::act_and_mul_quant_kernel<scalar_t, KERNEL<scalar_t>,  \
+                                             fp8_t>                       \
+                  <<<grid, block, 0, stream>>>(                           \
+                      out.mutable_data_ptr<fp8_t>(),                      \
+                      input.const_data_ptr<scalar_t>(),                   \
+                      scale.const_data_ptr<float>(), d);                  \
+            });                                                           \
      });

-void silu_and_mul_quant(torch::Tensor& out,    // [..., d]
-                        torch::Tensor& input,  // [..., 2 * d]
-                        torch::Tensor& scale) {
-  TORCH_CHECK(out.dtype() == torch::kFloat8_e4m3fn ||
-              out.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.dtype() == torch::kFloat16 ||
-              input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(input.size(-1) % 2 == 0);
+void silu_and_mul_quant(torch::stable::Tensor& out,    // [..., d]
+                        torch::stable::Tensor& input,  // [..., 2 * d]
+                        torch::stable::Tensor& scale) {
+  STD_TORCH_CHECK(
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(
+      input.scalar_type() == torch::headeronly::ScalarType::Half ||
+          input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Input must be FP16 or BF16");
+  STD_TORCH_CHECK(input.size(-1) % 2 == 0);
  LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel);
 }

 void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,              // (E, T, 2*H)
-    const at::Tensor& tokens_per_expert,  // (E)
-    at::Tensor& y_q,                      // (E, T, H) [OUT]
-    at::Tensor& y_s,                      // (E, T, H//group_size) [OUT]
+    const torch::stable::Tensor& input,              // (E, T, 2*H)
+    const torch::stable::Tensor& tokens_per_expert,  // (E)
+    torch::stable::Tensor& y_q,                      // (E, T, H) [OUT]
+    torch::stable::Tensor& y_s,  // (E, T, H//group_size) [OUT]
    bool cast_scale_ue8m0) {
 #ifndef USE_ROCM

@@ -606,14 +607,18 @@ void persistent_masked_m_silu_mul_quant(
  // fixed GROUP_SIZE of 128.
  static constexpr int GROUP_SIZE = 128;

-  TORCH_CHECK(input.dtype() == torch::kBFloat16);
-  TORCH_CHECK(y_q.dtype() == torch::kFloat8_e4m3fn ||
-              y_q.dtype() == torch::kFloat8_e4m3fnuz);
-  TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);
+  STD_TORCH_CHECK(input.scalar_type() ==
+                  torch::headeronly::ScalarType::BFloat16);
+  STD_TORCH_CHECK(
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
+      y_q.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz);
+  STD_TORCH_CHECK(input.size(-1) % (GROUP_SIZE * 2) == 0);

  bool const is_packed_ue8m0 =
-      (y_s.dtype() == torch::kInt32 && cast_scale_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kFloat32 || is_packed_ue8m0);
+      (y_s.scalar_type() == torch::headeronly::ScalarType::Int &&
+       cast_scale_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Float ||
+                  is_packed_ue8m0);

  using Idx_t = int64_t;

@@ -631,7 +636,7 @@ void persistent_masked_m_silu_mul_quant(

  int const NUM_GROUPS = H / GROUP_SIZE;

-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());

  // TODO: Get this from cuda_arch ?
  static constexpr int SILU_V2_BLOCK_COUNT = 132 * 32;
@@ -643,18 +648,21 @@ void persistent_masked_m_silu_mul_quant(
    static constexpr int max_shared_mem_bytes =                                \
        GROUP_SIZE * 2 * STAGES * NUM_WARPS * 2;                               \
    dim3 grid(sms), block(THREAD_COUNT);                                       \
-    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));          \
-    VLLM_DISPATCH_FP8_TYPES(                                                   \
+    const torch::stable::accelerator::DeviceGuard device_guard(                \
+        input.get_device_index());                                             \
+    VLLM_STABLE_DISPATCH_FP8_TYPES(                                            \
        y_q.scalar_type(), "silu_mul_fp8_quant_deep_gemm_kernel", [&] {        \
          vllm::silu_mul_fp8_quant_deep_gemm_kernel<                           \
              BLOCK_COUNT, max_shared_mem_bytes, fp8_t, scale_t, THREAD_COUNT, \
              Idx_t, CEIL_UE8M0, GROUP_SIZE, STAGES>                           \
              <<<grid, block, max_shared_mem_bytes + (E + 1) * 16, stream>>>(  \
-                  reinterpret_cast<__nv_bfloat16*>(input.data_ptr()),          \
-                  (fp8_t*)y_q.data_ptr(),                                      \
-                  reinterpret_cast<scale_t*>(y_s.data_ptr()),                  \
-                  reinterpret_cast<int32_t*>(tokens_per_expert.data_ptr()), E, \
-                  T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,       \
+                  reinterpret_cast<const __nv_bfloat16*>(                      \
+                      input.const_data_ptr()),                                 \
+                  y_q.mutable_data_ptr<fp8_t>(),                               \
+                  reinterpret_cast<scale_t*>(y_s.mutable_data_ptr()),          \
+                  reinterpret_cast<const int32_t*>(                            \
+                      tokens_per_expert.const_data_ptr()),                     \
+                  E, T, H, stride_i_e, stride_i_t, stride_i_h, stride_yq_e,    \
                  stride_yq_t, stride_yq_h, STRIDE_YS_E, STRIDE_YS_T,          \
                  STRIDE_YS_G, STRIDE_YS_P, stride_counts_e);                  \
        });
@@ -679,7 +687,7 @@ void persistent_masked_m_silu_mul_quant(
  Idx_t stride_ys_g = y_s.stride(2);
  Idx_t stride_ys_p = 0;
  if (!cast_scale_ue8m0) {
-    TORCH_CHECK(!is_packed_ue8m0);
+    STD_TORCH_CHECK(!is_packed_ue8m0);
    LAUNCH_ON_H(float, stride_ys_e, stride_ys_t, stride_ys_g, stride_ys_p,
                false);
    return;
@@ -692,8 +700,8 @@ void persistent_masked_m_silu_mul_quant(
    return;
  }

-  TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
-  TORCH_CHECK(y_s.dtype() == torch::kInt32);
+  STD_TORCH_CHECK(cast_scale_ue8m0 && is_packed_ue8m0);
+  STD_TORCH_CHECK(y_s.scalar_type() == torch::headeronly::ScalarType::Int);

  // Int32 packed ue8m0 scales tensor.
  // Let E, T, G be the number to experts, number of tokens and number of groups
@@ -31,7 +31,7 @@

 #include "cutlass/util/packed_stride.hpp"

-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"

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

 #include "cutlass/util/packed_stride.hpp"

-#include "core/math.hpp"
+#include "libtorch_stable/core/math.hpp"
 #include "core/batch_invariant.hpp"

 using namespace cute;
@@ -19,7 +19,7 @@
 #include "cutlass/gemm/collective/collective_builder.hpp"
 #include "cutlass/util/packed_stride.hpp"

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

@@ -14,7 +14,7 @@
 #include "cutlass/epilogue/collective/collective_builder.hpp"
 #include "cutlass/gemm/collective/collective_builder.hpp"

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

@@ -22,7 +22,7 @@
 #include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
 #include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"

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

@@ -301,8 +301,9 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(

  const int sf_k_local = local_group_id % kGroupsPerBlockX;
  const int row_local = local_group_id / kGroupsPerBlockX;
-  const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
-  const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
+  // Rows on grid.x: mn scales with tokens and can exceed the 65535 grid.y cap.
+  const int sf_k_idx = blockIdx.y * kGroupsPerBlockX + sf_k_local;
+  const int mn_idx = blockIdx.x * kRowsPerBlock + row_local;

 #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
  asm volatile("griddepcontrol.wait;");
@@ -496,14 +497,15 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
                  " is not a multiple of 4.");
  const int kx = GetGroupsPerBlockX(padded_groups_per_row);
  const int ry = 16 / kx;
-  const int64_t blocks_x = padded_groups_per_row / kx;
-  const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t row_blocks = (tma_aligned_mn + ry - 1) / ry;
+  const int64_t sf_k_blocks = padded_groups_per_row / kx;
  const int num_threads = (kx * ry) * THREADS_PER_GROUP;
-  // CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
-  STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
-                      blocks_y <= static_cast<int64_t>(INT32_MAX),
+  // CUDA caps grid.x at 2^31 - 1 and grid.y at 2^16 - 1 (65535).
+  constexpr int64_t kMaxGridDimYZ = 65535;
+  STD_TORCH_CHECK(row_blocks <= static_cast<int64_t>(INT32_MAX) &&
+                      sf_k_blocks <= kMaxGridDimYZ,
                  "per_token_group_quant_8bit_packed grid too large: (",
-                  blocks_x, ", ", blocks_y, ").");
+                  row_blocks, ", ", sf_k_blocks, ").");

  auto dst_type = output_q.scalar_type();

@@ -513,8 +515,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
    do {                                                                       \
      cudaLaunchConfig_t config = {};                                          \
-      config.gridDim = dim3(static_cast<unsigned int>(blocks_x),               \
-                            static_cast<unsigned int>(blocks_y));              \
+      config.gridDim = dim3(static_cast<unsigned int>(row_blocks),             \
+                            static_cast<unsigned int>(sf_k_blocks));           \
      config.blockDim = dim3(num_threads);                                     \
      config.dynamicSmemBytes = 0;                                             \
      config.stream = stream;                                                  \
@@ -539,8 +541,8 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
 #else
  #define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY)                         \
    do {                                                                       \
-      dim3 grid(static_cast<unsigned int>(blocks_x),                           \
-                static_cast<unsigned int>(blocks_y));                          \
+      dim3 grid(static_cast<unsigned int>(row_blocks),                         \
+                static_cast<unsigned int>(sf_k_blocks));                       \
      dim3 block(num_threads);                                                 \
      per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
                                                        RY>                    \
@@ -34,6 +34,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
  // TODO: Remove this once ROCm upgrade to torch 2.11.
  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");

+  // Note about marlin kernel 'workspace' arguments:
+  // Technically these should be mutable since they are modified by the kernel.
+  // But since they are set back to zero once the kernel is finished we can
+  // hand wave and say that they have no net effect.
+  //
+  // The reason to mark 'workspace' as immutable is so that they don't interfere
+  // with using ScalarType arguments in the ops. If they are marked as mutable,
+  // pytorch throws an assert in
+  // 'torch._higher_order_ops._register_effectful_op' that prevents these
+  // kernels from being torch.compile'd.
+  // See the following document for more info on custom types and ops that use
+  // custom types:
+  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
+
  // Machete (Dense) Optimized Mixed Precision GEMM for Hopper.
  ops.def(
      "machete_supported_schedules("
@@ -480,6 +494,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "Tensor workspace, int k, int max_seq_len) -> ()");

  // Activation ops
+  ops.def(
+      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
+      "y_q, Tensor! y_s, bool use_ue8m0) -> ()");
+  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
+
  // Activation function used in SwiGLU.
  ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");

@@ -492,6 +511,10 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
      "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit, "
      "float alpha=1.0, float beta=0.0) -> ()");

+  // SwiGLU activation with FP8 quantization.
+  ops.def(
+      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
+
  // Activation function used in GeGLU with `none` approximation.
  ops.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");

@@ -690,6 +713,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
  ops.impl("persistent_topk", TORCH_BOX(&persistent_topk));

  // Activation kernels (shared CUDA/ROCm)
+  ops.impl("persistent_masked_m_silu_mul_quant",
+           TORCH_BOX(&persistent_masked_m_silu_mul_quant));
+  ops.impl("weak_ref_tensor", TORCH_BOX(&weak_ref_tensor));
+  ops.impl("silu_and_mul_quant", TORCH_BOX(&silu_and_mul_quant));
  ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul));
  ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu));
  ops.impl("gelu_and_mul", TORCH_BOX(&gelu_and_mul));
@@ -9,28 +9,6 @@

 #include <vector>

-torch::Tensor weak_ref_tensor(torch::Tensor& tensor) {
-  // Ensure tensor is on CUDA
-  if (!tensor.is_cuda()) {
-    throw std::runtime_error("Tensor must be on CUDA device");
-  }
-
-  // Get the raw data pointer
-  void* data_ptr = tensor.data_ptr();
-
-  // Get tensor sizes and strides
-  std::vector<int64_t> sizes = tensor.sizes().vec();
-  std::vector<int64_t> strides = tensor.strides().vec();
-
-  // Get tensor options (dtype, device)
-  auto options = tensor.options();
-
-  // Create a new tensor from the raw data pointer
-  auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
-
-  return new_tensor;
-}
-
 // rms_norm and fused_add_rms_norm declarations also exist in
 // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here
 // because the CPU build still uses these torch::Tensor declarations.
@@ -53,16 +31,6 @@ void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
 void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit,
                        double alpha = 1.0, double beta = 0.0);

-void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
-                        torch::Tensor& scale);
-
-void persistent_masked_m_silu_mul_quant(
-    const at::Tensor& input,   // (E, T, 2*H)
-    const at::Tensor& counts,  // (E)
-    at::Tensor& y_q,           // (E, T, H) [OUT]
-    at::Tensor& y_s,           // (E, T, H//group_size) [OUT]
-    bool use_ue8m0);
-
 void gelu_and_mul(torch::Tensor& out, torch::Tensor& input);

 void gelu_tanh_and_mul(torch::Tensor& out, torch::Tensor& input);
@@ -0,0 +1,5 @@
+#include "core/registration.h"
+
+// QuTLASS registers torch.ops._qutlass_C via TORCH_LIBRARY in bindings.cpp.
+// This stub lets Python import vllm._qutlass_C to trigger op registration.
+REGISTER_EXTENSION(_qutlass_C)
@@ -20,17 +20,6 @@

 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // vLLM custom ops
-  //
-
-  ops.def(
-      "persistent_masked_m_silu_mul_quant(Tensor input, Tensor counts, Tensor! "
-      "y_q, Tensor! y_s,"
-      "bool use_ue8m0) -> ()");
-  ops.impl("persistent_masked_m_silu_mul_quant", torch::kCUDA,
-           &persistent_masked_m_silu_mul_quant);
-
-  ops.def("weak_ref_tensor(Tensor input) -> Tensor");
-  ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);

 #ifdef USE_ROCM
  // TODO: Remove this once we upgrade to torch 2.11.
@@ -39,35 +28,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
  ops.impl("get_cuda_view_from_cpu_tensor", torch::kCPU,
           &get_cuda_view_from_cpu_tensor);
-#endif
-
-  // Activation ops (quantized only — basic ops moved to _C_stable_libtorch)
-  ops.def(
-      "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
-  ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
-
-  // Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
-  // GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
-  // kernel launch. Registered in _C_stable_libtorch (incl. the FlashInfer V4
-  // full-cache bf16/fp8 variants).
-
-  // Quantization ops
-#ifndef USE_ROCM
-
-  // Note about marlin kernel 'workspace' arguments:
-  // Technically these should be mutable since they are modified by the kernel.
-  // But since they are set back to zero once the kernel is finished we can
-  // hand wave and say that they have no net effect.
-  //
-  // The reason to mark 'workspace' as immutable is so that they don't interfere
-  // with using ScalarType arguments in the ops. If they are marked as mutable,
-  // pytorch throws an assert in
-  // 'torch._higher_order_ops._register_effectful_op' that prevents these
-  // kernels from being torch.compile'd.
-  // See the following document for more info on custom types and ops that use
-  // custom types:
-  // https://docs.google.com/document/d/18fBMPuOJ0fY5ZQ6YyrHUppw9FA332CpNtgB6SOIgyuA
-
 #endif
 }

@@ -133,7 +133,7 @@ The model should inherit protocol `IsAttentionFree` and also implement class met
 For the mamba layers themselves, please use the [`MambaMixer`](../../../vllm/model_executor/layers/mamba/mamba_mixer.py) (for Mamba-1) or [`MambaMixer2`](../../../vllm/model_executor/layers/mamba/mamba_mixer2.py) (for Mamba-2) classes.
 The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/model_executor/models/config.py](../../../vllm/model_executor/models/config.py) to ensure that the runtime defaults are optimized.

-For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`BambaForCausalLM`](../../../vllm/model_executor/models/bamba.py) (for an example of a model that uses Mamba-2 and attention together).
+For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together).
 These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol).

 For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively.
@@ -40,7 +40,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test.
+schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
@@ -31,7 +31,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test.
+schemathesis>=4.0.0 # Required for openai schema test.
 # quantization
 bitsandbytes>=0.49.2
 buildkite-test-collector==0.1.9
@@ -39,7 +39,7 @@ lm-eval[api]>=0.4.12 # required for model evaluation test
 mteb[bm25s]>=2, <3 # required for mteb test
 transformers==5.5.3
 tokenizers==0.22.2
-schemathesis>=3.39.15 # Required for openai schema test
+schemathesis>=4.0.0 # Required for openai schema test
 # quantization
 bitsandbytes==0.49.2
 buildkite-test-collector==0.1.9
@@ -769,6 +769,7 @@ class precompiled_wheel_utils:
                            "vllm/_C.abi3.so",
                            "vllm/_C_stable_libtorch.abi3.so",
                            "vllm/_moe_C_stable_libtorch.abi3.so",
+                            "vllm/_qutlass_C.abi3.so",
                            "vllm/_flashmla_C.abi3.so",
                            "vllm/_flashmla_extension_C.abi3.so",
                            "vllm/_sparse_flashmla_C.abi3.so",
@@ -1135,6 +1136,7 @@ if _is_cuda():
        # DeepGEMM requires CUDA 12.3+ (SM90/SM100)
        # Optional since it won't build on unsupported architectures
        ext_modules.append(CMakeExtension(name="vllm._deep_gemm_C", optional=True))
+        ext_modules.append(CMakeExtension(name="vllm._qutlass_C", optional=True))
    # fmha_sm100 is a Python/CuTe-DSL package installed into vllm.third_party.
    ext_modules.append(CMakeExtension(name="vllm.fmha_sm100", optional=True))

@@ -1149,7 +1151,8 @@ if _is_cpu():
        ext_modules.append(CMakeExtension(name="vllm._C"))

 if _build_custom_ops():
-    ext_modules.append(CMakeExtension(name="vllm._C"))
+    if _is_hip():
+        ext_modules.append(CMakeExtension(name="vllm._C"))
    if _is_cuda() or _is_hip():
        ext_modules.append(CMakeExtension(name="vllm._C_stable_libtorch"))
        ext_modules.append(CMakeExtension(name="vllm._moe_C_stable_libtorch"))
@@ -22,7 +22,7 @@ import torch

 import vllm.config
 from tests.compile.backend import TestBackend
-from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
+from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
 from vllm.config import (
@@ -83,9 +83,8 @@ class _ViewDoubleQuantModel(torch.nn.Module):
    [_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
    ids=["no_view", "with_view"],
 )
-@pytest.mark.skipif(
-    not is_aiter_found_and_supported(),
-    reason="Only test on ROCm with AITER installed and supported",
+@pytest.mark.skip(
+    reason="Skipping for now because pytorch compiler removes one the two quant ops"
 )
 def test_double_aiter_rms_fp8_group_quant_fusion(
    model_cls: type[torch.nn.Module],
@@ -175,7 +175,7 @@ MULTIMODAL_MODELS = {
    "facebook/chameleon-7b": PPTestSettings.fast(),
    "adept/fuyu-8b": PPTestSettings.fast(),
    "zai-org/glm-4v-9b": PPTestSettings.fast(),
-    "OpenGVLab/InternVL2-1B": PPTestSettings.fast(),
+    "OpenGVLab/InternVL3-1B": PPTestSettings.fast(),
    "llava-hf/llava-1.5-7b-hf": PPTestSettings.fast(),
    "llava-hf/llava-v1.6-mistral-7b-hf": PPTestSettings.fast(),
    "llava-hf/LLaVA-NeXT-Video-7B-hf": PPTestSettings.fast(),
@@ -203,7 +203,7 @@ TEST_MODELS = [
    "intfloat/e5-mistral-7b-instruct",
    "BAAI/bge-multilingual-gemma2",
    # [MULTIMODAL GENERATION]
-    "OpenGVLab/InternVL2-1B",
+    "OpenGVLab/InternVL3-1B",
    "microsoft/Phi-3.5-vision-instruct",
    "fixie-ai/ultravox-v0_5-llama-3_2-1b",
    # [LANGUAGE GENERATION - HYBRID ARCH]
@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
    args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
    assert args.model == "s3://bucket/model"
    assert args.tokenizer == "s3://bucket/tokenizer"
+
+
+class TestDeviceIds:
+    def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
+        """--device-ids index beyond the CVD set raises ValueError."""
+        from vllm.platforms import current_platform
+
+        key = current_platform.device_control_env_var
+        monkeypatch.setenv(key, "4,5")
+        args = EngineArgs(model="m", device_ids=[0, 2])
+        with pytest.raises(ValueError, match="out of range"):
+            args._resolve_device_ids()
+
+    def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
+        """--device-ids are CVD-local indices resolved to physical ids."""
+        from vllm.platforms import current_platform
+
+        key = current_platform.device_control_env_var
+        monkeypatch.setenv(key, "4,5")
+        args = EngineArgs(model="m", device_ids=[0, 1])
+        assert args._resolve_device_ids() == [4, 5]
+
+    def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
+        """--device-ids support UUID CVD values resolved by the platform."""
+        from vllm.platforms import current_platform
+
+        key = current_platform.device_control_env_var
+        monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
+        monkeypatch.setattr(
+            type(current_platform),
+            "device_control_id_to_physical_device_id",
+            classmethod(
+                lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
+            ),
+        )
+
+        args = EngineArgs(model="m", device_ids=[0, 1])
+        assert args._resolve_device_ids() == [4, 5]
+
+    def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
+        """UUID --device-ids are resolved to physical IDs immediately."""
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(
+            type(current_platform),
+            "device_control_id_to_physical_device_id",
+            classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
+        )
+
+        args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
+        assert args._resolve_device_ids() == [4]
+
+    def test_device_ids_reject_mixed_integer_and_uuid_args(self):
+        """--device-ids must not mix CVD indices and UUIDs."""
+        args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
+        with pytest.raises(ValueError, match="must not mix"):
+            args._resolve_device_ids()
+
+    def test_no_device_ids(self):
+        """No --device-ids returns None."""
+        args = EngineArgs(model="m")
+        assert args._resolve_device_ids() is None
+
+    def test_cli_parsing(self):
+        """--device-ids parses comma-separated string from CLI."""
+        parser = FlexibleArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
+        assert parsed.device_ids == [0, 2, 4]
+
+    def test_cli_parsing_uuid(self):
+        """--device-ids parses comma-separated UUID strings from CLI."""
+        parser = FlexibleArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        parsed = parser.parse_args(
+            ["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
+        )
+        assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
+
+    def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
+        """assigned_physical_gpu_ids are already physical and not composed with CVD."""
+        import vllm.platforms.interface as platform_interface
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
+        monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
+
+        assert current_platform.device_id_to_physical_device_id(0) == 4
+        assert current_platform.device_id_to_physical_device_id(1) == 5
+        assert current_platform.logical_device_id_to_visible_device_id(0) == 0
+        assert current_platform.logical_device_id_to_visible_device_id(1) == 1
+
+    def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
+        """Physical IDs map back to visible ordinals when CVD uses UUIDs."""
+        import vllm.platforms.interface as platform_interface
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
+        monkeypatch.setenv(
+            current_platform.device_control_env_var,
+            "GPU-abcd1234,GPU-ef567890",
+        )
+        monkeypatch.setattr(
+            type(current_platform),
+            "device_control_id_to_physical_device_id",
+            classmethod(
+                lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
+            ),
+        )
+
+        assert current_platform.logical_device_id_to_visible_device_id(0) == 1
+
+    def test_device_ids_reject_duplicates(self):
+        """--device-ids must not contain duplicate entries."""
+        args = EngineArgs(model="m", device_ids=[2, 2])
+        with pytest.raises(ValueError, match="duplicates"):
+            args._resolve_device_ids()
+
+    def test_cli_parsing_strips_whitespace(self):
+        """--device-ids tolerates whitespace around commas."""
+        parser = FlexibleArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
+        assert parsed.device_ids == [0, 2, 4]
+
+    def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
+        """visible_device_id_to_physical_device_id maps torch device ordinals,
+        independent of the logical-to-physical mapping.
+
+        Regression test: CustomAllreduce passes device.index (a visible
+        ordinal) and must not index into assigned_physical_gpu_ids, which
+        raised IndexError for non-identity --device-ids like [2, 3].
+        """
+        import vllm.platforms.interface as platform_interface
+        from vllm.platforms import current_platform
+
+        monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
+        monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
+
+        # CVD unset: visible ordinal == physical ID, even beyond the
+        # assigned list's length.
+        assert current_platform.visible_device_id_to_physical_device_id(2) == 2
+        assert current_platform.visible_device_id_to_physical_device_id(3) == 3
+
+        monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
+        assert current_platform.visible_device_id_to_physical_device_id(1) == 5
+        with pytest.raises(IndexError, match="out of range"):
+            current_platform.visible_device_id_to_physical_device_id(2)
+
+
+class TestDpDeviceIdSharding:
+    def test_dp_supervisor_device_ids_stay_env_relative(self):
+        """Regression test: the DP supervisor must pass env-relative indices,
+        not physical IDs, because each child re-resolves --device-ids
+        against its inherited device-control env var."""
+        import argparse
+
+        from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
+
+        args = argparse.Namespace(
+            tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
+        )
+        assert _build_device_ids(args, local_rank=0) == [0, 1]
+        assert _build_device_ids(args, local_rank=1) == [2, 3]
+
+    def test_dp_supervisor_shards_user_device_ids(self):
+        """User-provided --device-ids are sharded across DP children."""
+        import argparse
+
+        from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
+
+        args = argparse.Namespace(
+            tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
+        )
+        assert _build_device_ids(args, local_rank=0) == [4, 5]
+        assert _build_device_ids(args, local_rank=1) == [6, 7]
+        with pytest.raises(ValueError, match="needs devices"):
+            _build_device_ids(args, local_rank=2)
+
+    def test_dp_rank_shards_user_assigned_gpu_ids(self):
+        """get_physical_gpu_ids_for_local_dp_rank slices the user-provided
+        --device-ids list instead of recomputing from the env var."""
+        from vllm.platforms import current_platform
+        from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
+
+        evar = current_platform.device_control_env_var
+        assert get_physical_gpu_ids_for_local_dp_rank(
+            evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
+        ) == [6, 7]
+        with pytest.raises(ValueError, match="needs devices"):
+            get_physical_gpu_ids_for_local_dp_rank(
+                evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
+            )
@@ -8,6 +8,8 @@ AnthropicServingMessages._convert_anthropic_to_openai_request().
 Also covers extended-thinking edge cases such as ``redacted_thinking``
 blocks echoed back by Anthropic clients, and streaming conversion in
 ``message_stream_converter``.
+
+Also covers cache usage computation in ``_build_anthropic_usage``.
 """

 import json
@@ -18,7 +20,11 @@ import pytest
 from vllm.entrypoints.anthropic.protocol import (
    AnthropicMessagesRequest,
 )
-from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
+from vllm.entrypoints.anthropic.serving import (
+    AnthropicServingMessages,
+    _build_anthropic_usage,
+    _get_cached_tokens,
+)
 from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionResponseStreamChoice,
    ChatCompletionStreamResponse,
@@ -27,6 +33,7 @@ from vllm.entrypoints.openai.engine.protocol import (
    DeltaFunctionCall,
    DeltaMessage,
    DeltaToolCall,
+    PromptTokenUsageInfo,
    UsageInfo,
 )

@@ -653,6 +660,108 @@ class TestThinkingBlockConversion:
        assert asst.get("content") == "Hi!"


+# ======================================================================
+# Cache usage computation
+# ======================================================================
+
+
+class TestGetCachedTokens:
+    """Tests for _get_cached_tokens helper."""
+
+    def test_none_usage(self):
+        assert _get_cached_tokens(None) is None
+
+    def test_no_prompt_tokens_details(self):
+        usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
+        assert _get_cached_tokens(usage) is None
+
+    def test_cached_tokens_present(self):
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+        )
+        assert _get_cached_tokens(usage) == 80
+
+    def test_cached_tokens_zero(self):
+        """Zero cached tokens should return 0, not None."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+        )
+        assert _get_cached_tokens(usage) == 0
+
+    def test_cached_tokens_none_in_details(self):
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
+        )
+        assert _get_cached_tokens(usage) is None
+
+
+class TestBuildAnthropicUsage:
+    """Tests for _build_anthropic_usage helper.
+
+    Anthropic defines: total_input = input_tokens + cache_read + cache_creation
+    vLLM's prompt_tokens is the total.
+    """
+
+    def test_no_cache_info(self):
+        """When cache info is unavailable, return raw prompt_tokens."""
+        result = _build_anthropic_usage(100, 10, None)
+        assert result.input_tokens == 100
+        assert result.output_tokens == 10
+        assert result.cache_read_input_tokens is None
+        assert result.cache_creation_input_tokens is None
+
+    def test_cache_hit(self):
+        """When cache is hit, input_tokens excludes cached tokens."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 20  # 100 - 80
+        assert result.output_tokens == 10
+        assert result.cache_read_input_tokens == 80
+        assert result.cache_creation_input_tokens == 0
+
+    def test_zero_cached_tokens(self):
+        """Zero cached tokens should still set cache_creation to 0."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 100  # 100 - 0
+        assert result.cache_read_input_tokens == 0
+        assert result.cache_creation_input_tokens == 0
+
+    def test_all_tokens_cached(self):
+        """When all tokens are cached, input_tokens should be 0."""
+        usage = UsageInfo(
+            prompt_tokens=100,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
+        )
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 0
+        assert result.cache_read_input_tokens == 100
+        assert result.cache_creation_input_tokens == 0
+
+    def test_no_prompt_tokens_details(self):
+        """UsageInfo without prompt_tokens_details returns no cache info."""
+        usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
+        result = _build_anthropic_usage(100, 10, usage)
+        assert result.input_tokens == 100
+        assert result.cache_read_input_tokens is None
+        assert result.cache_creation_input_tokens is None
+
+
 class TestInlineSystemMessageInMessagesArray:
    """Verify that ``role: system`` messages embedded inside the ``messages``
    array are preserved in their original position.
@@ -1098,6 +1207,135 @@ class TestMessageStartIncludesTypeAndRole:
        assert message["role"] == "assistant"


+class TestStreamingCacheUsageSemantics:
+    """Locks in the documented streaming behavior of cache usage fields.
+
+    vLLM's OpenAI chat completion streaming only attaches
+    ``prompt_tokens_details`` to the terminal usage chunk. The Anthropic layer
+    mirrors that contract: cache fields are omitted on ``message_start`` (key
+    absence signals "unknown") and populated on ``message_delta`` (the final
+    cumulative count). This is intentionally consistent with vLLM's OpenAI
+    behavior, even though Anthropic's upstream API populates cache fields on
+    ``message_start``; closing that gap requires plumbing cache info into the
+    first chunk at the OpenAI layer, which is out of scope here.
+    """
+
+    @pytest.mark.asyncio
+    async def test_streaming_cache_fields_absent_then_populated(self):
+        """First chunk lacks prompt_tokens_details (vLLM contract);
+        message_start omits cache fields. The final chunk carries
+        prompt_tokens_details, so message_delta carries resolved values."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant", content="hi"),
+                usage=UsageInfo(prompt_tokens=100, total_tokens=100),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(
+                    prompt_tokens=100,
+                    completion_tokens=5,
+                    total_tokens=105,
+                    prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
+                ),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        # message_start: cache fields unknown → omitted from JSON entirely.
+        start_usage = events[0][1]["message"]["usage"]
+        assert events[0][0] == "message_start"
+        assert start_usage["input_tokens"] == 100
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+
+        # message_delta: authoritative usage with cache fields populated.
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert delta_usage["input_tokens"] == 20  # 100 - 80
+        assert delta_usage["cache_read_input_tokens"] == 80
+        assert delta_usage["cache_creation_input_tokens"] == 0
+
+    @pytest.mark.asyncio
+    async def test_streaming_no_cache_hit(self):
+        """When the final chunk reports cached_tokens=0, message_delta carries
+        cache fields = 0 (cache miss); message_start still omits them."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=50, total_tokens=50),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(
+                    prompt_tokens=50,
+                    completion_tokens=5,
+                    total_tokens=55,
+                    prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
+                ),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        start_usage = events[0][1]["message"]["usage"]
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert start_usage["input_tokens"] == 50
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+        assert delta_usage["input_tokens"] == 50  # 50 - 0
+        assert delta_usage["cache_read_input_tokens"] == 0
+        assert delta_usage["cache_creation_input_tokens"] == 0
+
+    @pytest.mark.asyncio
+    async def test_streaming_no_prompt_tokens_details_at_all(self):
+        """If --enable-prompt-tokens-details is off, no chunk carries cache
+        info; both message_start and message_delta omit cache fields."""
+
+        async def sse_input():
+            yield _make_stream_chunk(
+                delta=DeltaMessage(role="assistant"),
+                usage=UsageInfo(prompt_tokens=30, total_tokens=30),
+            )
+            yield _make_stream_chunk(finish_reason="stop")
+            yield _make_stream_chunk(
+                choices=[],
+                usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32),
+            )
+            yield "data: [DONE]"
+
+        converter = _make_stream_converter()
+        output = []
+        async for event in converter.message_stream_converter(sse_input()):
+            output.append(event)
+        events = _parse_sse_events(output)
+
+        start_usage = events[0][1]["message"]["usage"]
+        delta_usage = next(
+            data["usage"] for ev, data in events if ev == "message_delta"
+        )
+        assert "cache_read_input_tokens" not in start_usage
+        assert "cache_creation_input_tokens" not in start_usage
+        assert "cache_read_input_tokens" not in delta_usage
+        assert "cache_creation_input_tokens" not in delta_usage
+
+
 # ======================================================================
 # Auto-detection of system-first template requirement
 # ======================================================================
@@ -364,7 +364,7 @@ class MockVLLMServer:
        await self._serve_task


-def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
+def launch_mock_vllm(child_args: argparse.Namespace):
    logger.info("Launching mock vLLM on port %s", child_args.port)
    mock_vllm = MockVLLMServer(
        port=child_args.port,
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]


 def launch_mock_vllm_with_drain(
-    child_args: argparse.Namespace, env_updates: dict[str, str]
+    child_args: argparse.Namespace,
 ):
    logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
    mock_vllm = MockVLLMServer(
@@ -6,15 +6,22 @@ from typing import Final
 import pytest
 import schemathesis
 from hypothesis import HealthCheck, settings
-from schemathesis import GenerationConfig
-from schemathesis.models import Case
+from schemathesis import GenerationMode
+from schemathesis.config import (
+    ChecksConfig,
+    CoveragePhaseConfig,
+    GenerationConfig,
+    PhasesConfig,
+    PositiveDataAcceptanceConfig,
+    ProjectConfig,
+    ProjectsConfig,
+    SchemathesisConfig,
+)

 from vllm.platforms import current_platform

 from ...utils import RemoteOpenAIServer

-schemathesis.experimental.OPEN_API_3_1.enable()
-
 MODEL_NAME = "HuggingFaceTB/SmolVLM-256M-Instruct"
 MAXIMUM_IMAGES = 2
 _ROCM_TIMEOUT_MULTIPLIER = 3 if current_platform.is_rocm() else 1
@@ -44,21 +51,38 @@ def server():
@pytest.fixture(scope="module")
 def get_schema(server):
    # avoid generating null (\x00) bytes in strings during test case generation
-    return schemathesis.openapi.from_uri(
+    return schemathesis.openapi.from_url(
        f"{server.url_root}/openapi.json",
-        generation_config=GenerationConfig(allow_x00=False),
+        config=SchemathesisConfig(
+            projects=ProjectsConfig(
+                default=ProjectConfig(
+                    generation=GenerationConfig(
+                        allow_x00=False,
+                        modes=[GenerationMode.POSITIVE],
+                    ),
+                    checks=ChecksConfig(
+                        positive_data_acceptance=PositiveDataAcceptanceConfig(
+                            enabled=False,
+                        ),
+                    ),
+                    phases=PhasesConfig(
+                        coverage=CoveragePhaseConfig(enabled=False),
+                    ),
+                ),
+            ),
+        ),
    )


-schema = schemathesis.from_pytest_fixture("get_schema")
+schema = schemathesis.pytest.from_fixture("get_schema")


@schemathesis.hook
-def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
+def before_generate_case(context: schemathesis.HookContext, strategy):
    op = context.operation
    assert op is not None

-    def no_invalid_types(case: schemathesis.models.Case):
+    def no_invalid_types(case: schemathesis.Case):
        """
        Skips tool_calls with `"type": "custom"` which schemathesis incorrectly
        generates instead of the valid `"type": "function"`.
@@ -68,39 +92,25 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
            -d '{"messages": [{"role": "assistant", "tool_calls": [{"custom": {"input": "", "name": ""}, "id": "", "type": "custom"}]}]}' \
            http://localhost:8000/v1/chat/completions
        """  # noqa: E501
-        if hasattr(case, "body") and isinstance(case.body, dict):
-            if (
-                "messages" in case.body
-                and isinstance(case.body["messages"], list)
-                and len(case.body["messages"]) > 0
-            ):
-                for message in case.body["messages"]:
-                    if not isinstance(message, dict):
-                        continue
+        if (
+            hasattr(case, "body")
+            and isinstance(case.body, dict)
+            and "messages" in case.body
+            and isinstance(case.body["messages"], list)
+            and len(case.body["messages"]) > 0
+        ):
+            for message in case.body["messages"]:
+                if not isinstance(message, dict):
+                    continue

-                    tool_calls = message.get("tool_calls", [])
-                    if isinstance(tool_calls, list):
-                        for tool_call in tool_calls:
-                            if isinstance(tool_call, dict):
-                                if tool_call.get("type") != "function":
-                                    return False
-                                if "custom" in tool_call:
-                                    return False
-
-            # Sometimes structured_outputs.grammar is generated to be empty
-            # Causing a server error in EBNF grammar parsing
-            # https://github.com/vllm-project/vllm/pull/22587#issuecomment-3195253421
-            structured_outputs = case.body.get("structured_outputs", {})
-            grammar = (
-                structured_outputs.get("grammar")
-                if isinstance(structured_outputs, dict)
-                else None
-            )
-
-            if grammar == "":
-                # Allow None (will be handled as no grammar)
-                # But skip empty strings
-                return False
+                tool_calls = message.get("tool_calls", [])
+                if isinstance(tool_calls, list):
+                    for tool_call in tool_calls:
+                        if isinstance(tool_call, dict):
+                            if tool_call.get("type") != "function":
+                                return False
+                            if "custom" in tool_call:
+                                return False

        return True

@@ -108,7 +118,6 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):


@schema.parametrize()
-@schema.override(headers={"Content-Type": "application/json"})
@settings(
    deadline=LONG_TIMEOUT_SECONDS * 1000,
    max_examples=50,
@@ -122,7 +131,7 @@ def before_generate_case(context: schemathesis.hooks.HookContext, strategy):
    # generating large-but-valid request bodies before vLLM is called.
    suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.data_too_large],
 )
-def test_openapi_stateless(case: Case):
+def test_openapi_stateless(case: schemathesis.Case):
    key = (
        case.operation.method.upper(),
        case.operation.path,
@@ -151,4 +160,8 @@ def test_openapi_stateless(case: Case):
    }.get(key, DEFAULT_TIMEOUT_SECONDS)

    # No need to verify SSL certificate for localhost
-    case.call_and_validate(verify=False, timeout=timeout)
+    case.call_and_validate(
+        verify=False,
+        timeout=timeout,
+        headers={"Content-Type": "application/json"},
+    )
@@ -25,7 +25,7 @@ def server():
        "--runner",
        "pooling",
        "--max-model-len",
-        "5000",
+        "16384",
        "--enforce-eager",
        "--limit-mm-per-prompt",
        json.dumps({"video": MAXIMUM_VIDEOS}),
@@ -143,4 +143,4 @@ def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
    assert output.model == model_name
    assert len(output.data) == 1
    assert len(output.data[0].probs) == 2
-    assert output.usage.prompt_tokens == 4807
+    assert output.usage.prompt_tokens == 8993
@@ -8,6 +8,7 @@ import pytest
 import pytest_asyncio

 from tests.utils import RemoteLaunchRenderServer
+from vllm.tokenizers import get_tokenizer

 MODEL_NAME = "hmellor/tiny-random-LlamaForCausalLM"

@@ -486,3 +487,438 @@ async def test_derender_completion_kv_transfer_params_passthrough(client):
    )
    assert response.status_code == 200
    assert response.json()["kv_transfer_params"] == kv
+
+
+# ---------------------------------------------------------------------------
+# E2E: render -> derender roundtrip with parser (reasoning + tool calls)
+# ---------------------------------------------------------------------------
+
+PARSER_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+
+_E2E_TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_weather",
+            "description": "Get weather for a city",
+            "parameters": {
+                "type": "object",
+                "properties": {"city": {"type": "string"}},
+            },
+        },
+    }
+]
+
+
+@pytest.fixture(scope="module")
+def parser_server():
+    args = [
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "hermes",
+        "--reasoning-parser",
+        "deepseek_r1",
+    ]
+    with RemoteLaunchRenderServer(PARSER_MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def parser_client(parser_server):
+    async with httpx.AsyncClient(
+        base_url=parser_server.url_for(""), timeout=60.0
+    ) as http_client:
+        yield http_client
+
+
+@pytest.fixture(scope="module")
+def parser_tokenizer():
+    return get_tokenizer(PARSER_MODEL)
+
+
+def _encode(tokenizer, text: str) -> list[int]:
+    return tokenizer.encode(text, add_special_tokens=False)
+
+
+def _decoded(tokenizer, token_ids: list[int]) -> str:
+    return tokenizer.decode(token_ids, skip_special_tokens=True)
+
+
+def _require_markers_survive(tokenizer, text: str, *markers: str) -> list[int]:
+    """Encode text and skip the test if any marker is lost in roundtrip."""
+    ids = _encode(tokenizer, text)
+    decoded = tokenizer.decode(ids, skip_special_tokens=False)
+    for m in markers:
+        if m not in decoded:
+            pytest.skip(f"Marker {m!r} lost in encode->decode roundtrip")
+    return ids
+
+
+async def _e2e_render_chat(
+    client: httpx.AsyncClient,
+    model: str,
+    messages: list[dict],
+) -> dict:
+    resp = await client.post(
+        "/v1/chat/completions/render",
+        json={"model": model, "messages": messages},
+    )
+    assert resp.status_code == 200, resp.text
+    return resp.json()
+
+
+def _e2e_generate_response(
+    token_ids: list[int],
+    request_id: str = "chatcmpl-e2e-test",
+) -> dict:
+    return {
+        "request_id": request_id,
+        "choices": [
+            {
+                "index": 0,
+                "token_ids": token_ids,
+                "finish_reason": "stop",
+            }
+        ],
+    }
+
+
+@pytest.mark.asyncio
+async def test_e2e_plain_roundtrip(parser_client, parser_tokenizer):
+    """Plain text without reasoning markers roundtrips correctly."""
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "The answer is four."
+    output_ids = _encode(parser_tokenizer, answer)
+    expected = _decoded(parser_tokenizer, output_ids)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert content == expected
+
+
+@pytest.mark.asyncio
+async def test_e2e_token_identity(parser_client, parser_tokenizer):
+    """encode(derender(token_ids)) == token_ids (RL invariant)."""
+    messages = [{"role": "user", "content": "Hi"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "Hello! How can I help?"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    re_encoded = _encode(parser_tokenizer, content)
+    assert output_ids == re_encoded
+
+
+@pytest.mark.asyncio
+async def test_e2e_non_ascii_roundtrip(parser_client, parser_tokenizer):
+    """CJK + emoji roundtrip without U+FFFD."""
+    messages = [{"role": "user", "content": "Reply in Chinese"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "你好世界 😀"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert "�" not in content
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer):
+    """<think>...</think> splits into reasoning + content."""
+    messages = [{"role": "user", "content": "What is 2+3?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    reasoning_text = "The user wants 2 plus 3. That is 5."
+    answer_text = "The answer is 5."
+    output_text = f"<think>{reasoning_text}</think>{answer_text}"
+    output_ids = _require_markers_survive(parser_tokenizer, output_text, "</think>")
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    msg = resp.json()["choices"][0]["message"]
+    assert msg["reasoning"] is not None
+    assert reasoning_text in msg["reasoning"]
+    assert answer_text in msg["content"]
+    assert "<think>" not in msg["content"]
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer):
+    """<tool_call> extracted into tool_calls field."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    output_text = (
+        "<think>Let me check the weather.</think>"
+        '<tool_call>\n{"name": "get_weather", '
+        '"arguments": {"city": "Paris"}}\n</tool_call>'
+    )
+    output_ids = _require_markers_survive(
+        parser_tokenizer,
+        output_text,
+        "</think>",
+        "<tool_call>",
+        "</tool_call>",
+    )
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "tools": _E2E_TOOLS,
+                "tool_choice": "auto",
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    choice = resp.json()["choices"][0]
+    assert choice["message"]["tool_calls"]
+    assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather"
+
+
+@pytest.mark.asyncio
+async def test_e2e_parsed_reasoning_and_tool_call(parser_client, parser_tokenizer):
+    """Reasoning + tool call in the same output."""
+    messages = [{"role": "user", "content": "Weather in Paris?"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    reasoning_text = "I should look up the weather."
+    tool_text = (
+        '<tool_call>\n{"name": "get_weather", '
+        '"arguments": {"city": "Paris"}}\n</tool_call>'
+    )
+    output_text = f"<think>{reasoning_text}</think>{tool_text}"
+    output_ids = _require_markers_survive(
+        parser_tokenizer, output_text, "</think>", "<tool_call>"
+    )
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": PARSER_MODEL,
+                "messages": messages,
+                "tools": _E2E_TOOLS,
+                "tool_choice": "auto",
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    choice = resp.json()["choices"][0]
+    assert choice["message"]["reasoning"] is not None
+    assert reasoning_text in choice["message"]["reasoning"]
+    assert choice["message"]["tool_calls"]
+
+
+@pytest.mark.asyncio
+async def test_e2e_no_chat_request_fallback(parser_client, parser_tokenizer):
+    """Without chat_request, derender falls back to plain detokenization."""
+    messages = [{"role": "user", "content": "Hello"}]
+    gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages)
+
+    answer = "Hi there!"
+    output_ids = _encode(parser_tokenizer, answer)
+
+    resp = await parser_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": PARSER_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+        },
+    )
+    assert resp.status_code == 200
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert "Hi" in content
+
+
+# ---------------------------------------------------------------------------
+# E2E: HarmonyParser + GPT-OSS
+# ---------------------------------------------------------------------------
+
+HARMONY_MODEL = "openai/gpt-oss-20b"
+
+
+def _ensure_harmony_vocab():
+    """Pre-cache the o200k_base BPE file needed by openai-harmony.
+
+    The Rust tiktoken-rs backend downloads from Azure Blob Storage, which
+    may be unreachable in some environments.  When the cache is cold we
+    fetch the file ourselves and place it in ``/tmp/tiktoken-rs-cache/``
+    using the SHA-1(URL) filename that tiktoken-rs expects.
+    """
+    import hashlib
+    import urllib.request
+    from pathlib import Path
+
+    url = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
+    cache_dir = Path("/tmp/tiktoken-rs-cache")
+    cache_key = hashlib.sha1(url.encode()).hexdigest()
+    cache_file = cache_dir / cache_key
+    if not cache_file.exists():
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, cache_file)
+
+
+@pytest.fixture(scope="module")
+def harmony_server():
+    _ensure_harmony_vocab()
+    args = [
+        "--trust-remote-code",
+        "--enable-auto-tool-choice",
+        "--tool-call-parser",
+        "openai",
+        "--reasoning-parser",
+        "openai_gptoss",
+    ]
+    with RemoteLaunchRenderServer(HARMONY_MODEL, args) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def harmony_client(harmony_server):
+    async with httpx.AsyncClient(
+        base_url=harmony_server.url_for(""), timeout=60.0
+    ) as http_client:
+        yield http_client
+
+
+@pytest.fixture(scope="module")
+def harmony_tokenizer():
+    return get_tokenizer(HARMONY_MODEL, trust_remote_code=True)
+
+
+def _harmony_extract_assistant_ids(
+    tokenizer, assistant_msg: dict, user_content: str = "test"
+) -> list[int]:
+    """Extract assistant token IDs via apply_chat_template diff."""
+    prompt = [{"role": "user", "content": user_content}]
+    full = prompt + [assistant_msg]
+    text_prompt = tokenizer.apply_chat_template(
+        prompt, add_generation_prompt=True, tokenize=False
+    )
+    text_full = tokenizer.apply_chat_template(
+        full, add_generation_prompt=False, tokenize=False
+    )
+    prompt_ids = tokenizer.encode(text_prompt)
+    full_ids = tokenizer.encode(text_full)
+    assistant_ids = list(full_ids[len(prompt_ids) :])
+    if not assistant_ids:
+        pytest.skip("Could not extract assistant tokens for Harmony")
+    return assistant_ids
+
+
+@pytest.mark.asyncio
+async def test_e2e_harmony_plain_roundtrip(harmony_client, harmony_tokenizer):
+    """GPT-OSS content-only roundtrip."""
+    messages = [{"role": "user", "content": "What is 2+2?"}]
+    gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
+
+    assistant_msg = {"role": "assistant", "content": "Four."}
+    output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
+
+    resp = await harmony_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": HARMONY_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": HARMONY_MODEL,
+                "messages": messages,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    content = resp.json()["choices"][0]["message"]["content"]
+    assert content is not None and len(content) > 0
+    assert "Four" in content
+
+
+@pytest.mark.asyncio
+async def test_e2e_harmony_reasoning(harmony_client, harmony_tokenizer):
+    """GPT-OSS reasoning: analysis channel extracted."""
+    messages = [{"role": "user", "content": "Add 2 and 3."}]
+    gen_req = await _e2e_render_chat(harmony_client, HARMONY_MODEL, messages)
+
+    reasoning_text = "The user wants 2 plus 3."
+    answer_text = "The answer is 5."
+    assistant_msg = {
+        "role": "assistant",
+        "thinking": reasoning_text,
+        "content": answer_text,
+    }
+    output_ids = _harmony_extract_assistant_ids(harmony_tokenizer, assistant_msg)
+
+    decoded = harmony_tokenizer.decode(output_ids)
+    if reasoning_text not in decoded:
+        pytest.skip("Harmony template did not render thinking")
+
+    resp = await harmony_client.post(
+        "/v1/chat/completions/derender",
+        json={
+            "model": HARMONY_MODEL,
+            "generate_response": _e2e_generate_response(output_ids),
+            "prompt_tokens": len(gen_req["token_ids"]),
+            "chat_request": {
+                "model": HARMONY_MODEL,
+                "messages": messages,
+                "include_reasoning": True,
+            },
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    msg = resp.json()["choices"][0]["message"]
+    assert msg["reasoning"] is not None
+    assert reasoning_text in msg["reasoning"]
+    assert answer_text in (msg["content"] or "")
@@ -78,7 +78,16 @@ def test_gsm8k_correctness(config_filename):
            "Skipping DeepSeek-V3.2 and DeepSeek-R1 on ROCm platforms "
            "due to agent pool disk space issues and pod evictions."
        )
+    if current_platform.is_rocm() and (
+        "Qwen3.5-35B-A3B-MXFP4-AITER-TP2" in config_filename.name
+    ):
+        from vllm.platforms.rocm import on_gfx950

+        if not on_gfx950():
+            pytest.skip(
+                "Skipping Qwen3.5-35B-A3B-MXFP4-AITER-TP2 on non-GFX950 platforms. "
+                "The quantization scheme is not supported on non-GFX950 platforms."
+            )
    # Parse server arguments from config (use shlex to handle quoted strings)
    server_args_str = eval_config.get("server_args", "")
    server_args = shlex.split(server_args_str) if server_args_str else []
@@ -15,16 +15,14 @@ from vllm.config import (
 from vllm.platforms import current_platform
 from vllm.platforms.cpu import CpuPlatform

-# CudaPlatform and RocmPlatform import their respective compiled C extensions
-# at module level, raising ModuleNotFoundError on incompatible builds.
-try:
+if current_platform.is_cuda():
    from vllm.platforms.cuda import CudaPlatform
-except (ImportError, ModuleNotFoundError):
+else:
    CudaPlatform = None

-try:
+if current_platform.is_rocm():
    from vllm.platforms.rocm import RocmPlatform
-except (ImportError, ModuleNotFoundError):
+else:
    RocmPlatform = None

 from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -434,9 +432,15 @@ def test_per_head_quant_scales_backend_selection(
    [
        ("FLASH_ATTN", True, True),  # FlashAttn supports non-causal
        ("FLASH_ATTN", False, True),  # FlashAttn also works with causal
-        ("FLASHINFER", True, False),  # FlashInfer does not support non-causal
-        ("FLASHINFER", False, True),  # FlashInfer works with causal
-    ],
+    ]
+    + (
+        [
+            ("FLASHINFER", True, False),  # FlashInfer does not support non-causal
+            ("FLASHINFER", False, True),  # FlashInfer works with causal
+        ]
+        if CudaPlatform is not None
+        else []
+    ),
 )
 def test_non_causal_backend_selection(
    backend_name: str, use_non_causal: bool, should_succeed: bool
@@ -459,11 +463,12 @@ def test_non_causal_backend_selection(
        attention_config=attention_config, cache_config=cache_config
    )

-    if CudaPlatform is None:
-        pytest.skip("CudaPlatform not available")
+    platform = CudaPlatform or RocmPlatform
+    if platform is None:
+        pytest.skip("CudaPlatform and RocmPlatform are not available")
    with (
        set_current_vllm_config(vllm_config),
-        patch("vllm.platforms.current_platform", CudaPlatform()),
+        patch("vllm.platforms.current_platform", platform()),
    ):
        if should_succeed:
            backend = get_attn_backend(
@@ -5,10 +5,12 @@ import math
 import random
 import time
 from collections.abc import Callable
+from contextlib import nullcontext

 import pytest
 import torch
 import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel

 from vllm.platforms import current_platform
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, set_random_seed
@@ -557,15 +559,21 @@ def test_contexted_kv_attention_alibi(
            query_len, seq_len, alibi_slopes, device, dtype
        )

-        # Compute attention
-        out = F.scaled_dot_product_attention(
-            q_sdpa,
-            k_sdpa,
-            v_sdpa,
-            attn_mask=alibi_mask,
-            dropout_p=0.0,
-            scale=scale,
-        )
+        # Compute attention. On ROCm we force use of the Math SDPA backend rather than
+        # the Flash or Mem-Efficient backends for increased numerical accuracy
+        if current_platform.is_rocm():
+            sdpa_context = sdpa_kernel(SDPBackend.MATH)
+        else:
+            sdpa_context = nullcontext()
+        with sdpa_context:
+            out = F.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=alibi_mask,
+                dropout_p=0.0,
+                scale=scale,
+            )

        # Reshape output back to [query_len, num_heads, head_size]
        out = out.view(num_heads, query_len, head_size).permute(1, 0, 2)
@@ -90,7 +90,9 @@ def _ref_sparse_prefill_ragged(
    return out.to(torch.bfloat16)


-def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
+def _pack_fp8_ds_mla_cache(
+    kv: torch.Tensor, block_size: int, is_extra: bool = False
+) -> torch.Tensor:
    assert kv.shape[-1] == HEAD_DIM
    num_tokens = kv.shape[0]
    num_blocks = (num_tokens + block_size - 1) // block_size
@@ -101,7 +103,9 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
    )
    cache_flat = cache.view(torch.uint8).flatten()
    kv_nope_fp8 = (
-        kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
+        kv[:, :NOPE_HEAD_DIM]
+        .to(torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype())
+        .view(torch.uint8)
    )
    kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)

@@ -120,7 +124,7 @@ def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:


 def _read_fp8_ds_mla_cache(
-    cache: torch.Tensor, slot: int, block_size: int
+    cache: torch.Tensor, slot: int, block_size: int, is_extra: bool = False
 ) -> torch.Tensor:
    cache_flat = cache.view(torch.uint8).flatten()
    block_idx = slot // block_size
@@ -129,7 +133,9 @@ def _read_fp8_ds_mla_cache(
    token_base = block_base + pos * 576

    nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
-    nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
+    nope = nope_u8.view(
+        torch.float8_e4m3fn if is_extra else current_platform.fp8_dtype()
+    ).to(torch.float32)
    rope_u8 = cache_flat[
        token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
    ]
@@ -157,7 +163,9 @@ def _ref_sparse_decode_ragged(
        ]
        if extra_cache is not None and extra_rows is not None:
            row_kv.extend(
-                _read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
+                _read_fp8_ds_mla_cache(
+                    extra_cache, int(slot), block_size, is_extra=True
+                )
                for slot in extra_rows[query_idx]
            )

@@ -326,7 +334,7 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
    main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
    extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
    main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
-    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+    extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
    main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
    main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
    extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
@@ -477,7 +485,7 @@ def test_sparse_attn_decode_split_k_kernel(
        rows = [[1, 3, 0, 5, 2, 4], [3, 0, 6]]
        extra_kv = torch.randn(7, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
        extra_rows = rows
-        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
+        extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size, is_extra=True)
        extra_indices, extra_indptr = _ragged_from_rows(rows, device)

    attn_sink = (
@@ -18,11 +18,7 @@ HEAD_SIZES = [128, 256]
 BLOCK_SIZES = [16]

 DTYPES = [torch.bfloat16]
-QDTYPES = (
-    [None, torch.float8_e4m3fn]
-    if not current_platform.is_rocm()
-    else [None, torch.float8_e4m3fnuz]
-)
+QDTYPES = [None, current_platform.fp8_dtype()]
 FP8_DTYPE = current_platform.fp8_dtype()

 # one value large enough to test overflow in index calculation.
@@ -10,8 +10,12 @@ from torch.multiprocessing import spawn
 from tests.kernels.utils import opcheck
 from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
 from vllm.distributed import cleanup_dist_env_and_memory
-from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
+from vllm.model_executor.layers.minimax_rms_norm import (
+    MiniMaxText01RMSNormTP,
+    rms_norm_tp,
+)
 from vllm.platforms import current_platform
+from vllm.triton_utils import HAS_TRITON
 from vllm.utils.network_utils import get_open_port
 from vllm.utils.torch_utils import set_random_seed

@@ -54,8 +58,19 @@ def _worker_forward_qk(
    torch.manual_seed(seed + 1000 + local_rank)
    qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")

-    q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
-    ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
+    # Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces
+    # the variance (it is the tp==1 / already-reduced building block), so the
+    # multi-rank reference must use the eager path that performs the global
+    # variance all-reduce, matching the fused kernel below.
+    ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
+        qkv.clone(),
+        q_norm.weight,
+        k_norm.weight,
+        hq,
+        hk,
+        world_size,
+        eps,
+    )

    # Set up Lamport workspace.
    from vllm.distributed.parallel_state import get_tp_group
@@ -150,3 +165,44 @@ def test_minimax_reduce_rms_qk(
        nprocs=world_size,
        join=True,
    )
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda() or not HAS_TRITON,
+    reason="CUDA and Triton required",
+)
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049])
+@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("tp_world", [1, 4, 8])
+@pytest.mark.parametrize("eps", [1e-6])
+@pytest.mark.parametrize("seed", [42])
+def test_minimax_qk_norm_triton_fallback(
+    monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed
+):
+    """Single-GPU check: Triton fallback kernels vs the pure-torch reference.
+
+    The all-reduce is a TP communication barrier, so it is monkeypatched to
+    identity here; both the Triton path and the reference see the same
+    (patched) reduction. This validates the kernel math and the folded
+    ``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
+    are the per-rank q/k segment widths.
+    """
+    monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
+
+    q_size, kv_size = hidden_dims
+    device = "cuda"
+    torch.manual_seed(seed)
+    qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device)
+    q_weight = torch.randn(q_size, dtype=dtype, device=device)
+    k_weight = torch.randn(kv_size, dtype=dtype, device=device)
+
+    q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback(
+        qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps
+    )
+    q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager(
+        qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
+    )
+
+    torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2)
+    torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2)
@@ -9,6 +9,7 @@ import pytest
 import torch
 from packaging import version

+from vllm._aiter_ops import is_aiter_found
 from vllm.platforms import current_platform
 from vllm.utils.flashinfer import has_flashinfer

@@ -31,17 +32,15 @@ HOPPER_MXFP4_BF16_AVAILABLE = (
 # ROCm platform and dependencies
 ROCM_AVAILABLE = current_platform.is_rocm()
 ROCM_TRITON_KERNELS_AVAILABLE = False
-ROCM_AITER_AVAILABLE = False
+ROCM_AITER_AVAILABLE = is_aiter_found()
 ROCM_GFX950 = False

 if ROCM_AVAILABLE:
-    from vllm._aiter_ops import rocm_aiter_ops
    from vllm.platforms.rocm import on_gfx950
    from vllm.utils.import_utils import has_triton_kernels

    ROCM_TRITON_KERNELS_AVAILABLE = has_triton_kernels()
    ROCM_GFX950 = on_gfx950()
-    ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_enabled()

    if ROCM_AITER_AVAILABLE:
        from aiter.ops.triton.moe.quant_moe import upcast_from_mxfp
@@ -83,7 +82,7 @@ def enable_pickle(monkeypatch):
    [
        ModelCase("fxmarty/qwen_1.5-moe-a2.7b-mxfp4", tp=2),
        ModelCase("fxmarty/deepseek_r1_3_layers_mxfp4", tp=8),
-        ModelCase("fxmarty/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
+        ModelCase("mawong-amd/Llama-4-Scout-17B-16E-Instruct-2-layers-mxfp4", tp=1),
        ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=1),
        ModelCase("fxmarty/Llama-3.1-70B-Instruct-2-layers-mxfp6", tp=4),
    ],
@@ -102,6 +101,7 @@ def test_mxfp4_loading_and_execution_moe(vllm_runner, model_case: ModelCase):
        tensor_parallel_size=model_case.tp,
        load_format="dummy",
        compilation_config={"cudagraph_capture_sizes": [16]},
+        gpu_memory_utilization=0.8,  # mxfp6 models use more scratch space
    ) as llm:
        # Disabled as check_model is broken: https://github.com/vllm-project/vllm/pull/18465#issuecomment-3329880562
        # def check_model(model):
@@ -1267,7 +1267,7 @@ def test_rocm_mxfp4_moe_oracle(

    This test validates that the oracle functions work end-to-end:
    - select_mxfp4_moe_backend() selects a valid backend
-    - convert_to_mxfp4_moe_kernel_format() converts weights without error
+    - convert_gpt_oss_weight_to_mxfp4_moe_kernel_format() converts weights without error
    - make_mxfp4_moe_quant_config() builds a valid quant config
    - make_mxfp4_moe_kernel() creates a kernel that runs without error
    - The kernel output is within accuracy tolerance of reference
@@ -1287,7 +1287,7 @@ def test_rocm_mxfp4_moe_oracle(
    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
        Mxfp4MoeBackend,
        backend_to_kernel_cls,
-        convert_to_mxfp4_moe_kernel_format,
+        convert_gpt_oss_weight_to_mxfp4_moe_kernel_format,
        make_mxfp4_moe_kernel,
        make_mxfp4_moe_quant_config,
    )
@@ -1387,7 +1387,7 @@ def test_rocm_mxfp4_moe_oracle(

    # Convert weights using oracle
    w13_conv, w2_conv, w13_scale_conv, w2_scale_conv, w13_bias_conv, w2_bias_conv = (
-        convert_to_mxfp4_moe_kernel_format(
+        convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(
            mxfp4_backend=backend,
            layer=layer,  # type: ignore[arg-type]
            w13_weight=w13_quant,
@@ -1423,7 +1423,7 @@ def test_rocm_mxfp4_moe_oracle(
            mxfp4_backend=backend,
            experts_cls=experts_cls,
            routing_tables=None,
-            shared_experts=None,
+            layer=None,
        )

        # Create inputs
@@ -345,6 +345,63 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
        )


+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
+)
+def test_per_token_group_quant_fp8_packed_large_mn():
+    """Regression test for https://github.com/vllm-project/vllm/issues/45099.
+
+    Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
+    2^16 - 1, respectively.
+    Prior code introduced a bug where it incorrectly assumed grid.x and y both have
+    2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
+    launched with a large mn that exceeds grid.y limit (2^16 - 1).
+
+    This issue doesn't surface often because each forward pass only processes a
+    bounded token batch, not the full context.
+    Quantizing tensors with more rows than that will fail at launch with
+    "CUDA error: invalid argument".
+    This is a differential test that compares fp8 output against Triton output
+    reference when token size sits just above the gridDim.y 2^16 - 1 limit.
+    """
+
+    device = "cuda"
+    group_size = 128
+    # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
+    # row, so any mn > 65535 overflowed grid.y before the fix.
+    num_tokens, hidden_dim = 65537, 2048
+    torch.manual_seed(42)
+    x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
+
+    out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
+        x,
+        group_size=group_size,
+        use_ue8m0=True,
+    )
+
+    with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
+        ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
+            x, group_size, use_ue8m0=True
+        )
+
+    assert torch.equal(out_q, ref_q), "Quantized output mismatch"
+
+    # Vectorized packed-scale check; the per-element loop used by the smaller
+    # tests is too slow at this size. groups_per_row is a multiple of 4 here,
+    # so there is no K padding and the packed view lines up.
+    mn = num_tokens
+    groups_per_row = hidden_dim // group_size
+    k_num_packed = (groups_per_row + 3) // 4
+    assert groups_per_row % 4 == 0
+    ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
+    exp = ref_exponents.view(mn, k_num_packed, 4)
+    expected = (
+        exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
+    )
+    assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
+
+
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@@ -60,8 +60,10 @@ def test_rocm_compressed_tensors_w8a8(
    vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
 ):
    dtype = "bfloat16"
-
-    with vllm_runner(model_path, dtype=dtype) as vllm_model:
+    # Pin to TRITON_ATTN, see https://github.com/vllm-project/vllm/issues/46179
+    with vllm_runner(
+        model_path, dtype=dtype, attention_backend="TRITON_ATTN"
+    ) as vllm_model:
        vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)


@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 import pytest
+import torch
 import torch.nn.functional as F
 import transformers.utils
 from PIL import Image
@@ -52,6 +53,7 @@ def _get_cherry_blossom_image() -> Image.Image:
    )


+@torch.inference_mode()
 def _run_test(
    hf_runner: type[HfRunner],
    vllm_runner: type[VllmRunner],
@@ -92,3 +92,95 @@ def test_processor_num_frames_timestamp(
    assert len(video_phs) == 1, (
        f"Expected exactly 1 video placeholder, got {len(video_phs)}"
    )
+
+
+@pytest.mark.parametrize("model_id", [MODEL_ID])
+@pytest.mark.parametrize("num_videos", [2, 4])
+def test_processor_multi_video(
+    model_id: str,
+    num_videos: int,
+) -> None:
+    """Verify that multi-video processing produces correct placeholders.
+
+    This exercises the token-level replacement path in
+    ``_call_hf_processor`` which avoids the quadratic text-level
+    prompt expansion.
+    """
+    ctx = build_model_context(
+        model_id,
+        limit_mm_per_prompt={"image": 0, "video": num_videos},
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+
+    prompt = "<|vision_start|><|video_pad|><|vision_end|>" * num_videos
+    mm_data = {"video": [_build_video_mm_data(num_frames=8)["video"][0]] * num_videos}
+
+    processed = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data(mm_data),
+        hf_processor_mm_kwargs={"num_frames": 8},
+    )
+
+    token_ids = processed["prompt_token_ids"]
+    assert len(token_ids) > 0
+
+    video_phs = processed["mm_placeholders"].get("video", [])
+    assert len(video_phs) == num_videos, (
+        f"Expected {num_videos} video placeholders, got {len(video_phs)}"
+    )
+
+    # All placeholders should have the same length (same video params)
+    # and must not overlap.
+    lengths = {ph.length for ph in video_phs}
+    assert len(lengths) == 1, f"Placeholder lengths differ: {lengths}"
+    for i in range(1, len(video_phs)):
+        prev_end = video_phs[i - 1].offset + video_phs[i - 1].length
+        assert video_phs[i].offset >= prev_end, (
+            f"Placeholder {i} overlaps with placeholder {i - 1}"
+        )
+
+
+@pytest.mark.parametrize("model_id", [MODEL_ID])
+@pytest.mark.parametrize(
+    "hf_mm_kwargs",
+    [{"num_frames": [8, 16]}, {"fps": [2.0, 4.0]}],
+)
+def test_processor_multi_video_list_kwargs(
+    model_id: str,
+    hf_mm_kwargs: dict[str, Any],
+) -> None:
+    """Regression test: a multi-video request with list-valued per-video
+    ``mm_processor_kwargs`` (one ``fps``/``num_frames`` per video) must not
+    crash.
+
+    Before the fix, ``_call_hf_processor`` copied the whole kwargs to every
+    video without slicing, so ``_get_video_second_idx`` received the list
+    where a scalar was expected and raised ``TypeError``.
+    """
+    ctx = build_model_context(
+        model_id,
+        limit_mm_per_prompt={"image": 0, "video": 2},
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+
+    prompt = (
+        "<|vision_start|><|video_pad|><|vision_end|>"
+        "<|vision_start|><|video_pad|><|vision_end|>"
+    )
+    mm_data = {
+        "video": [
+            _build_video_mm_data(num_frames=16)["video"][0],
+            _build_video_mm_data(num_frames=32)["video"][0],
+        ]
+    }
+
+    processed = processor(
+        prompt,
+        mm_items=processor.info.parse_mm_data(mm_data),
+        hf_processor_mm_kwargs=hf_mm_kwargs,
+    )
+
+    video_phs = processed["mm_placeholders"].get("video", [])
+    assert len(video_phs) == 2, (
+        f"Expected exactly 2 video placeholders, got {len(video_phs)}"
+    )
@@ -1530,6 +1530,16 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
        "Qwen/Qwen3-VL-8B-Instruct",
        speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
    ),
+    "Eagle3Qwen3ForCausalLM": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        trust_remote_code=True,
+        speculative_model=(
+            "inference-optimization/"
+            "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1"
+        ),
+        tokenizer="Qwen/Qwen3-8B",
+        use_original_num_layers=True,
+    ),
    # [PEagle]
    "PEagleDraftModel": _HfExamplesInfo(
        "Qwen/Qwen3-8B",
@@ -1545,6 +1555,15 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
        tokenizer="Qwen/Qwen3-8B",
        use_original_num_layers=True,
    ),
+    "PeagleQwen3ForCausalLM": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        trust_remote_code=True,
+        speculative_model=(
+            "inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4"
+        ),
+        tokenizer="Qwen/Qwen3-8B",
+        use_original_num_layers=True,
+    ),
    # [MTP]
    "DeepSeekMTPModel": _HfExamplesInfo(
        "luccafong/deepseek_mtp_main_random",
@@ -46,7 +46,8 @@ def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float():

 def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
    vllm_config = SimpleNamespace(
-        scheduler_config=SimpleNamespace(max_num_batched_tokens=4)
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
+        compilation_config=SimpleNamespace(static_forward_context={}),
    )
    experts = DeepseekV4MegaMoEExperts(
        vllm_config,
@@ -182,3 +183,81 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
        fused_topk_weights.view(torch.uint8),
        ref_topk_weights.view(torch.uint8),
    )
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
+)
+def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding():
+    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
+
+    device = torch.device("cuda")
+    num_tokens = 7
+    hidden_size = 256
+    top_k = 8
+
+    generator = torch.Generator(device=device)
+    generator.manual_seed(1)
+    hidden_states = torch.randn(
+        num_tokens,
+        hidden_size,
+        device=device,
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    topk_ids = torch.randint(
+        0,
+        256,
+        (num_tokens, top_k),
+        device=device,
+        dtype=torch.int32,
+        generator=generator,
+    )
+    topk_weights = torch.randn(
+        num_tokens,
+        top_k,
+        device=device,
+        dtype=torch.float32,
+        generator=generator,
+    )
+    is_padding = torch.tensor(
+        [False, True, False, False, True, False, True],
+        device=device,
+    )
+
+    ref_x, ref_x_sf = per_token_cast_to_fp8(
+        hidden_states,
+        use_ue8m0=True,
+        gran_k=32,
+        use_packed_ue8m0=True,
+    )
+    ref_topk_idx = topk_ids.to(torch.int64)
+    ref_topk_idx[is_padding] = -1
+    ref_topk_weights = topk_weights.clone()
+    ref_topk_weights[is_padding] = 0.0
+
+    fused_x = torch.empty_like(ref_x)
+    fused_x_sf = torch.empty_like(ref_x_sf)
+    fused_topk_idx = torch.empty_like(ref_topk_idx)
+    fused_topk_weights = torch.empty_like(ref_topk_weights)
+
+    prepare_megamoe_inputs(
+        hidden_states,
+        topk_weights,
+        topk_ids,
+        fused_x,
+        fused_x_sf,
+        fused_topk_idx,
+        fused_topk_weights,
+        is_padding=is_padding,
+    )
+    torch.accelerator.synchronize()
+
+    assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
+    assert torch.equal(fused_x_sf, ref_x_sf)
+    assert torch.equal(fused_topk_idx, ref_topk_idx)
+    assert torch.equal(
+        fused_topk_weights.view(torch.uint8),
+        ref_topk_weights.view(torch.uint8),
+    )
@@ -15,6 +15,7 @@ from vllm.multimodal.video import (
    DynamicVideoBackend,
    GLM46VVideoBackend,
    Molmo2VideoBackend,
+    Qwen2VLVideoBackend,
    Qwen3VLVideoBackend,
    VideoLoader,
    VideoSourceMetadata,
@@ -70,11 +71,12 @@ def test_video_loader_type_doesnt_exist():


@pytest.mark.parametrize(
-    "model_repo, expected_loader_cls",
+    "model_repo, expected_loader_cls, hf_sample_kwargs",
    [
        pytest.param(
            "allenai/Molmo2-4B",
            Molmo2VideoBackend,
+            None,
            marks=pytest.mark.skip(
                reason="Video processor not aligned, investigate later.",
            ),
@@ -83,23 +85,44 @@ def test_video_loader_type_doesnt_exist():
        pytest.param(
            "zai-org/GLM-4.1V-9B-Thinking",
            DynamicVideoBackend,
+            None,
            id="glm4v",
        ),
        pytest.param(
            "zai-org/GLM-4.6V-Flash",
            GLM46VVideoBackend,
+            None,
            id="glm46v",
        ),
        pytest.param(
            "Qwen/Qwen3-VL-4B-Instruct",
            Qwen3VLVideoBackend,
+            None,
            id="qwen3vl",
        ),
+        # Qwen2-VL/Qwen2.5-VL ship no ``video_processor_type`` in their
+        # preprocessor config, so resolution relies on the model_type ->
+        # video processor fallback in get_video_processor_cls_name_from_config.
+        # They also ship no default fps/num_frames, so the HF sampler needs an
+        # explicit target rate; pass fps=2 to match the loader default.
+        pytest.param(
+            "Qwen/Qwen2-VL-7B-Instruct",
+            Qwen2VLVideoBackend,
+            {"fps": 2},
+            id="qwen2vl",
+        ),
+        pytest.param(
+            "Qwen/Qwen2.5-VL-7B-Instruct",
+            Qwen2VLVideoBackend,
+            {"fps": 2},
+            id="qwen2_5_vl",
+        ),
    ],
 )
 def test_video_processor_from_model_repo(
    model_repo: str,
    expected_loader_cls: type,
+    hf_sample_kwargs: dict[str, int | float] | None,
 ):
    """Test that a model repo resolves to the correct video loader backend.

@@ -143,7 +166,7 @@ def test_video_processor_from_model_repo(
            fps=vllm_meta["fps"],
            duration=vllm_meta["duration"],
        )
-        hf_indices = processor.sample_frames(hf_metadata)
+        hf_indices = processor.sample_frames(hf_metadata, **(hf_sample_kwargs or {}))
        vllm_indices = np.array(vllm_meta["frames_indices"])
        np.testing.assert_array_equal(
            hf_indices,
@@ -96,6 +96,9 @@ class MockTokenizer:
        return "".join(parts)


+CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
+
+
 def make_mock_tokenizer(sample: Sample) -> MockTokenizer:
    """Build a mock tokenizer from a sample's vocab and token data."""
    return MockTokenizer(
@@ -19,6 +19,7 @@ import pytest
 from pydantic import TypeAdapter

 from tests.parser.engine.replay_harness import (
+    CHUNK_SIZES,
    MockTokenizer,
    assert_parse_output,
    collect_output,
@@ -113,8 +114,6 @@ _PAIRINGS = _discover_pairings()

 _ALL_SAMPLES = [(p.parser_cls, s) for p in _PAIRINGS for s in p.samples]

-CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
-

@pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
@pytest.mark.parametrize(
@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression test for U+FFFD leak at reasoning→content transition.
+
+When byte-fallback tokens span the reasoning/content boundary,
+decoding isolated content-side token IDs via tokenizer.decode()
+produces U+FFFD (Unicode replacement character).  The fix flushes
+the reasoning parser's engine lexer instead.
+
+Reproduces the bug at various chunk sizes and validates that the
+fix prevents U+FFFD from leaking into streamed content.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tests.parser.engine.replay_harness import (
+    CHUNK_SIZES,
+    MockTokenizer,
+    collect_output,
+    replay_streaming,
+)
+from vllm.parser.abstract_parser import DelegatingParser
+from vllm.parser.engine.registered_adapters import (
+    Glm47MoeParserReasoningAdapter,
+    Glm47MoeParserToolAdapter,
+    Qwen3ParserReasoningAdapter,
+    Qwen3ParserToolAdapter,
+)
+
+
+class ByteFallbackMockTokenizer(MockTokenizer):
+    """MockTokenizer that returns U+FFFD for specified token IDs.
+
+    Simulates byte-fallback tokenizer behavior where isolated
+    partial-byte tokens decode to the Unicode replacement character.
+    """
+
+    def __init__(
+        self,
+        vocab: dict[str, int],
+        tokens: list[tuple[int, str]],
+        ufffd_token_ids: set[int],
+    ) -> None:
+        super().__init__(vocab, tokens)
+        self._ufffd_token_ids = frozenset(ufffd_token_ids)
+
+    def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
+        parts: list[str] = []
+        for tid in ids:
+            if skip_special_tokens and tid in self._special_ids:
+                continue
+            if tid in self._ufffd_token_ids:
+                parts.append("�")
+            else:
+                text = self._token_decode_map.get(tid, f"?{tid}?")
+                parts.append(text)
+        return "".join(parts)
+
+
+# ── Model-specific DelegatingParser subclasses ───────────────────────
+
+
+class _Glm47Delegating(DelegatingParser):
+    reasoning_parser_cls = Glm47MoeParserReasoningAdapter
+    tool_parser_cls = Glm47MoeParserToolAdapter
+
+
+class _Qwen3Delegating(DelegatingParser):
+    reasoning_parser_cls = Qwen3ParserReasoningAdapter
+    tool_parser_cls = Qwen3ParserToolAdapter
+
+
+# ── Shared test data ─────────────────────────────────────────────────
+
+_SHARED_TOKENS: list[tuple[int, str]] = [
+    (100, "Let me"),
+    (101, " think"),
+    (102, " about"),
+    (103, " Samsung."),
+    (51, "</think>"),
+    (200, "삼성"),
+    (201, "전자의"),
+    (202, " 주가를"),
+    (203, " 분석합니다."),
+]
+
+_SHARED_UFFFD_IDS: set[int] = {200}
+
+EXPECTED_REASONING = "Let me think about Samsung."
+EXPECTED_CONTENT = "삼성전자의 주가를 분석합니다."
+
+_MODEL_CONFIGS = [
+    pytest.param(
+        {
+            "<think>": 50,
+            "</think>": 51,
+            "<tool_call>": 60,
+            "</tool_call>": 61,
+            "<arg_key>": 62,
+            "</arg_key>": 63,
+            "<arg_value>": 64,
+            "</arg_value>": 65,
+        },
+        _Glm47Delegating,
+        id="glm47",
+    ),
+    pytest.param(
+        {
+            "<think>": 50,
+            "</think>": 51,
+            "<tool_call>": 60,
+            "</tool_call>": 61,
+        },
+        _Qwen3Delegating,
+        id="qwen3",
+    ),
+]
+
+
+# ── Tests ────────────────────────────────────────────────────────────
+
+
+class TestUfffdReasoningTransition:
+    """U+FFFD must not appear at the reasoning→content transition."""
+
+    @pytest.mark.parametrize("vocab,delegating_cls", _MODEL_CONFIGS)
+    @pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
+    def test_no_ufffd(self, chunk_size, vocab, delegating_cls):
+        tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
+        parser = delegating_cls(tokenizer)
+        deltas = replay_streaming(
+            parser,
+            _SHARED_TOKENS,
+            chunk_size=chunk_size,
+            finished_on_last=True,
+        )
+        output = collect_output(deltas)
+
+        assert "�" not in output.content, (
+            f"U+FFFD leaked into content: {output.content!r}"
+        )
+        assert output.content == EXPECTED_CONTENT
+        assert output.reasoning == EXPECTED_REASONING
+
+    def test_byte_fallback_tokenizer_produces_ufffd(self):
+        """Validate the fixture: decode() returns U+FFFD for isolated
+        byte-fallback token IDs, proving the old code path would leak."""
+        vocab = dict(_MODEL_CONFIGS[0].values[0])
+        tokenizer = ByteFallbackMockTokenizer(vocab, _SHARED_TOKENS, _SHARED_UFFFD_IDS)
+        assert tokenizer.decode([200]) == "�"
+
+    @pytest.mark.parametrize("chunk_size", CHUNK_SIZES, ids=lambda c: f"chunk={c}")
+    def test_multiple_ufffd_tokens_at_boundary(self, chunk_size):
+        """Multiple consecutive byte-fallback tokens at the boundary."""
+        tokens: list[tuple[int, str]] = [
+            (100, "Reasoning."),
+            (51, "</think>"),
+            (200, "삼"),
+            (201, "성"),
+            (202, "전자"),
+        ]
+        ufffd_ids: set[int] = {200, 201}
+        vocab = dict(_MODEL_CONFIGS[0].values[0])
+
+        tokenizer = ByteFallbackMockTokenizer(vocab, tokens, ufffd_ids)
+        parser = _Glm47Delegating(tokenizer)
+        deltas = replay_streaming(
+            parser,
+            tokens,
+            chunk_size=chunk_size,
+            finished_on_last=True,
+        )
+        output = collect_output(deltas)
+
+        assert "�" not in output.content, (
+            f"U+FFFD leaked into content: {output.content!r}"
+        )
+        assert output.content == "삼성전자"
+        assert output.reasoning == "Reasoning."
@@ -61,8 +61,8 @@ class QuantConfig:
    quant_max: float
    quant_min: float
    kv_quant_mode: KVQuantMode
-    # INT8 Triton stores truncate; FP8 hardware casts round.
-    uses_trunc: bool
+    # INT8 rounds explicitly; FP8 relies on dtype cast rounding.
+    rounds_before_store: bool


 INT8_CONFIG = QuantConfig(
@@ -71,7 +71,7 @@ INT8_CONFIG = QuantConfig(
    quant_max=127.0,
    quant_min=-128.0,
    kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
-    uses_trunc=True,
+    rounds_before_store=True,
 )
 FP8_CONFIG = QuantConfig(
    cache_dtype=FP8_DTYPE,
@@ -79,7 +79,7 @@ FP8_CONFIG = QuantConfig(
    quant_max=FP8_MAX,
    quant_min=FP8_MIN,
    kv_quant_mode=KVQuantMode.FP8_PER_TOKEN_HEAD,
-    uses_trunc=False,
+    rounds_before_store=False,
 )

 QUANT_CONFIGS = [INT8_CONFIG, FP8_CONFIG]
@@ -104,7 +104,7 @@ def _quantize_per_token_head_ref(
    absmax = data.float().abs().amax(dim=2)  # [num_tokens, num_heads]
    scales = (absmax / cfg.quant_max).clamp(min=1e-6)
    scaled = data.float() * (1.0 / scales[:, :, None])
-    if cfg.uses_trunc:
+    if cfg.rounds_before_store:
        q = scaled.round().clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
    else:
        q = scaled.clamp(cfg.quant_min, cfg.quant_max).to(cfg.cache_dtype)
@@ -255,7 +255,7 @@ def test_per_token_head_round_trip_accuracy(
 ):
    """Verify per-token-head round-trip: kernel dequant matches reference.

-    INT8: Triton truncates on float->int8 store.
+    INT8: round-to-nearest before int8 store.
    FP8: hardware cast (clamp then cast).
    """
    from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
@@ -315,6 +315,52 @@ def test_per_token_head_round_trip_accuracy(
                )


+@torch.inference_mode()
+def test_int8_per_token_head_raw_cache_matches_round_reference():
+    """INT8 cache writes should match round-to-nearest quantization exactly."""
+    from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
+        triton_reshape_and_cache_flash_per_token_head_quant,
+    )
+
+    torch.set_default_device(DEVICE_TYPE)
+
+    head_size = 8
+    block_size = 4
+
+    key = torch.tensor(
+        [[[-127.0, -2.6, -2.4, -1.6, -1.4, -0.6, -0.4, 127.0]]],
+        dtype=torch.bfloat16,
+    )
+    value = -key
+
+    key_cache = torch.zeros(1, block_size, 1, head_size, dtype=torch.int8)
+    value_cache = torch.zeros_like(key_cache)
+    k_scale_cache = torch.ones(1, block_size, 1, dtype=torch.float32)
+    v_scale_cache = torch.ones_like(k_scale_cache)
+    slot_mapping = torch.tensor([2], dtype=torch.long)
+
+    triton_reshape_and_cache_flash_per_token_head_quant(
+        key,
+        value,
+        key_cache,
+        value_cache,
+        k_scale_cache,
+        v_scale_cache,
+        slot_mapping,
+    )
+
+    ref_k_quant, ref_k_scales = _quantize_per_token_head_ref(key, INT8_CONFIG)
+    ref_v_quant, ref_v_scales = _quantize_per_token_head_ref(value, INT8_CONFIG)
+
+    slot = slot_mapping.item()
+    blk = slot // block_size
+    off = slot % block_size
+    assert torch.equal(key_cache[blk, off], ref_k_quant[0])
+    assert torch.equal(value_cache[blk, off], ref_v_quant[0])
+    torch.testing.assert_close(k_scale_cache[blk, off], ref_k_scales[0])
+    torch.testing.assert_close(v_scale_cache[blk, off], ref_v_scales[0])
+
+
 # ===========================================================================
 # 4. Negative slot mapping (padding tokens should be skipped)
 # ===========================================================================
@@ -461,7 +507,7 @@ def test_triton_unified_attention_per_token_head_scale(

    scaled_k = key_cache_bf16.float() / k_scale_cache[:, :, :, None]
    scaled_v = value_cache_bf16.float() / v_scale_cache[:, :, :, None]
-    if qcfg.uses_trunc:
+    if qcfg.rounds_before_store:
        key_cache_q = (
            scaled_k.round().clamp(qcfg.quant_min, qcfg.quant_max).to(qcfg.cache_dtype)
        )
@@ -18,6 +18,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
    FP8ScaledMMLinearLayerConfig,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    get_fp8_min_max,
    kFp8DynamicTokenSym,
    kFp8StaticChannelSym,
    kFp8StaticTensorSym,
@@ -309,7 +310,7 @@ def test_hipb_mm_kernel_forward_accuracy(enable_hipb_mm_kernel):
    _check_bpreshuffle_runtime_support(weight_shape, num_tokens=num_tokens)

    fp8_dtype = current_platform.fp8_dtype()
-    fp8_max = torch.finfo(fp8_dtype).max
+    fp8_max = get_fp8_min_max()[1]
    device = torch.device("cuda")

    # Build a bf16 weight and quantize per output channel (one scale per row).
@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for contiguous KV cache packing."""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm import envs
+from vllm.v1.core.kv_cache_utils import (
+    _get_kv_cache_config_deepseek_v4,
+    get_kv_cache_config_from_groups,
+)
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+    MLAAttentionSpec,
+    SlidingWindowSpec,
+    UniformTypeKVCacheSpecs,
+)
+
+
+def _make_mla_spec(page_size: int, block_size: int = 256) -> MLAAttentionSpec:
+    return MLAAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=512,
+        dtype=torch.uint8,
+        page_size_padded=page_size,
+        cache_dtype_str="fp8_ds_mla",
+        model_version="deepseek_v4",
+        alignment=576,
+    )
+
+
+def _make_full_spec() -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.float16,
+    )
+
+
+def _make_sw_spec() -> SlidingWindowSpec:
+    return SlidingWindowSpec(
+        block_size=16,
+        num_kv_heads=2,
+        head_size=64,
+        dtype=torch.float16,
+        sliding_window=128,
+    )
+
+
+def _make_groups(n_c4, n_c128, n_swa):
+    PS_C4_MLA = 37440
+    PS_C4_IDX = 8640
+    PS_C128 = 1728
+    PS_SWA = 37440
+
+    mla_specs = {}
+    for i in range(n_c4):
+        mla_specs[f"c4_mla.{i}"] = _make_mla_spec(PS_C4_MLA)
+        mla_specs[f"c4_idx.{i}"] = _make_mla_spec(PS_C4_IDX)
+    for i in range(n_c128):
+        mla_specs[f"c128_mla.{i}"] = _make_mla_spec(PS_C128)
+
+    mla_group = KVCacheGroupSpec(
+        layer_names=list(mla_specs.keys()),
+        kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=mla_specs),
+    )
+
+    swa_specs = {}
+    for i in range(n_swa):
+        swa_specs[f"swa.{i}"] = _make_mla_spec(PS_SWA)
+
+    swa_group = KVCacheGroupSpec(
+        layer_names=list(swa_specs.keys()),
+        kv_cache_spec=UniformTypeKVCacheSpecs(block_size=256, kv_cache_specs=swa_specs),
+    )
+
+    return [mla_group, swa_group]
+
+
+def _mock_vllm_config():
+    config = MagicMock()
+    config.cache_config.num_gpu_blocks_override = None
+    return config
+
+
+def _run(n_c4=3, n_c128=2, n_swa=5, mem=100 * 1024 * 1024):
+    groups = _make_groups(n_c4, n_c128, n_swa)
+    return _get_kv_cache_config_deepseek_v4(_mock_vllm_config(), groups, mem)
+
+
+def _page_sizes_by_layer(
+    groups: list[KVCacheGroupSpec],
+) -> dict[str, int]:
+    page_sizes = {}
+    for group in groups:
+        specs = group.kv_cache_spec.kv_cache_specs
+        for layer_name in group.layer_names:
+            page_sizes[layer_name] = specs[layer_name].page_size_bytes
+    return page_sizes
+
+
+class TestInterleavedPacking:
+    def test_all_tensors_have_block_stride(self):
+        _, tensors = _run()
+        for t in tensors:
+            assert t.block_stride > 0
+
+    def test_all_tensors_share_same_size(self):
+        _, tensors = _run()
+        sizes = set(t.size for t in tensors)
+        assert len(sizes) == 1
+        assert sizes.pop() > 0
+
+    def test_offsets_within_one_block(self):
+        _, tensors = _run()
+        for t in tensors:
+            assert t.offset < t.block_stride
+
+    def test_all_layers_accounted_for(self):
+        n_c4, n_c128, n_swa = 5, 4, 7
+        _, tensors = _run(n_c4=n_c4, n_c128=n_c128, n_swa=n_swa)
+        all_names = set()
+        for t in tensors:
+            all_names.update(t.shared_by)
+        expected = n_c4 * 2 + n_c128 + n_swa
+        assert len(all_names) == expected
+
+    def test_strided_views_are_independent(self):
+        groups = _make_groups(n_c4=3, n_c128=2, n_swa=5)
+        page_sizes = _page_sizes_by_layer(groups)
+        num_blocks, tensors = _get_kv_cache_config_deepseek_v4(
+            _mock_vllm_config(), groups, 100 * 1024 * 1024
+        )
+        backing = torch.zeros(tensors[0].size, dtype=torch.uint8)
+        views = []
+        for t in tensors:
+            page_size = page_sizes[t.shared_by[0]]
+            v = torch.as_strided(
+                backing,
+                size=(num_blocks, page_size),
+                stride=(t.block_stride, 1),
+                storage_offset=t.offset,
+            )
+            views.append(v)
+
+        for i, v in enumerate(views):
+            v.fill_(i + 1)
+
+        for i, v in enumerate(views):
+            assert (v == i + 1).all(), f"View {i} was corrupted"
+
+    def test_hma_attention_groups_keep_default_backing(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", False, raising=False)
+        full = _make_full_spec()
+        sw = _make_sw_spec()
+        page_size = full.page_size_bytes
+        groups = [
+            KVCacheGroupSpec(["full.0", "full.1"], full),
+            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
+            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
+        ]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+        )
+
+        assert config.num_blocks == 32
+        assert sum(t.size for t in config.kv_cache_tensors) == page_size * 2 * 32
+        assert config.kv_cache_tensors == [
+            KVCacheTensor(size=page_size * 32, shared_by=["full.0", "sw.0", "sw.1"]),
+            KVCacheTensor(size=page_size * 32, shared_by=["full.1", "sw.2", "sw.3"]),
+        ]
+
+    def test_hma_attention_groups_use_packed_backing_with_flag(self, monkeypatch):
+        monkeypatch.setattr(envs, "VLLM_USE_PACKED_HMA_KV_CACHE", True, raising=False)
+        full = _make_full_spec()
+        sw = _make_sw_spec()
+        page_size = full.page_size_bytes
+        groups = [
+            KVCacheGroupSpec(["full.0", "full.1"], full),
+            KVCacheGroupSpec(["sw.0", "sw.2"], sw),
+            KVCacheGroupSpec(["sw.1", "sw.3"], sw),
+        ]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=page_size * 2 * 32
+        )
+
+        assert config.num_blocks == 32
+        assert {t.size for t in config.kv_cache_tensors} == {page_size * 2 * 32}
+        assert config.kv_cache_tensors == [
+            KVCacheTensor(
+                size=page_size * 2 * 32,
+                shared_by=["full.0", "sw.0", "sw.1"],
+                offset=0,
+                block_stride=page_size * 2,
+            ),
+            KVCacheTensor(
+                size=page_size * 2 * 32,
+                shared_by=["full.1", "sw.2", "sw.3"],
+                offset=page_size,
+                block_stride=page_size * 2,
+            ),
+        ]
+
+    def test_single_group_attention_keeps_unpacked_layout(self):
+        spec = _make_full_spec()
+        groups = [KVCacheGroupSpec(["full.0", "full.1"], spec)]
+
+        config = get_kv_cache_config_from_groups(
+            _mock_vllm_config(), groups, available_memory=spec.page_size_bytes * 2 * 32
+        )
+
+        assert sum(t.size for t in config.kv_cache_tensors) == (
+            spec.page_size_bytes * 2 * 32
+        )
+        assert [t.block_stride for t in config.kv_cache_tensors] == [0, 0]
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
@@ -117,6 +117,7 @@ def new_kv_cache_spec(
    page_size_padded=None,
    sliding_window=None,
    attention_chunk_size=None,
+    indexes_kv_by_block_stride=False,
 ):
    return FullAttentionSpec(
        block_size=block_size,
@@ -126,6 +127,7 @@ def new_kv_cache_spec(
        page_size_padded=page_size_padded,
        sliding_window=sliding_window,
        attention_chunk_size=attention_chunk_size,
+        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
    )


@@ -136,6 +138,7 @@ def new_sliding_window_spec(
    dtype=torch.float32,
    page_size_padded=None,
    sliding_window=1,
+    indexes_kv_by_block_stride=False,
 ):
    return SlidingWindowSpec(
        block_size=block_size,
@@ -144,6 +147,7 @@ def new_sliding_window_spec(
        dtype=dtype,
        page_size_padded=page_size_padded,
        sliding_window=sliding_window,
+        indexes_kv_by_block_stride=indexes_kv_by_block_stride,
    )


@@ -1799,16 +1803,38 @@ def test_get_kv_cache_config_one_worker():
        ],
    )

-    # different hidden size that cannot be aligned by using different block size
+    # different hidden size that cannot be aligned by using different block size,
+    # but can be aligned by padding the smaller physical page.
+    swa_spec = new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True)
    kv_cache_specs_hybrid = {
-        "layer_1": new_kv_cache_spec(head_size=64),
-        "layer_2": new_sliding_window_spec(head_size=96),
+        "layer_1": new_kv_cache_spec(head_size=64, indexes_kv_by_block_stride=True),
+        "layer_2": swa_spec,
    }

-    with pytest.raises(NotImplementedError):
-        get_kv_cache_configs(
-            vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
-        )[0]
+    kv_cache_config_hybrid = get_kv_cache_configs(
+        vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 2 * 32]
+    )[0]
+    padded_page_size = swa_spec.page_size_bytes
+    assert kv_cache_config_hybrid == KVCacheConfig(
+        num_blocks=42,
+        kv_cache_tensors=[
+            KVCacheTensor(size=padded_page_size * 42, shared_by=["layer_1", "layer_2"]),
+        ],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["layer_1"],
+                new_kv_cache_spec(
+                    head_size=64,
+                    page_size_padded=padded_page_size,
+                    indexes_kv_by_block_stride=True,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["layer_2"],
+                new_sliding_window_spec(head_size=96, indexes_kv_by_block_stride=True),
+            ),
+        ],
+    )

    # Test num_gpu_blocks_override
    vllm_config.cache_config.num_gpu_blocks_override = 16
@@ -2322,6 +2348,75 @@ def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
        get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])


+def test_unify_kv_cache_page_size_uses_padding_for_non_divisible_sizes():
+    """DFlash drafters can have a smaller head size than the target model.
+
+    For example, MiMo uses 192-dim target KV heads while its DFlash draft uses
+    128-dim KV heads. The resulting page sizes are 3:2 rather than an integer
+    block-size multiple, so the smaller page must be padded instead.
+    """
+    # Both layers' backends opt into the padded-page strided view (e.g.
+    # FlashAttention / its DiffKV subclass), so padding is allowed.
+    target_spec = new_kv_cache_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=192,
+        dtype=torch.bfloat16,
+        indexes_kv_by_block_stride=True,
+    )
+    draft_spec = new_sliding_window_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+        sliding_window=1024,
+        indexes_kv_by_block_stride=True,
+    )
+
+    unified_specs = kv_cache_utils.unify_kv_cache_spec_page_size(
+        {
+            "target_attn": target_spec,
+            "draft_attn": draft_spec,
+        }
+    )
+
+    assert unified_specs["target_attn"] == target_spec
+    unified_draft_spec = unified_specs["draft_attn"]
+    assert unified_draft_spec.block_size == draft_spec.block_size
+    assert unified_draft_spec.real_page_size_bytes == draft_spec.real_page_size_bytes
+    assert unified_draft_spec.page_size_padded == target_spec.page_size_bytes
+    assert unified_draft_spec.page_size_bytes == target_spec.page_size_bytes
+
+
+def test_unify_kv_cache_page_size_padding_requires_backend_support():
+    """Padding is gated on the backend declaring ``indexes_kv_by_block_stride``.
+
+    A backend that does not support the strided padded-page view must raise
+    rather than silently padding (and misreading KV at runtime).
+    """
+    target_spec = new_kv_cache_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=192,
+        dtype=torch.bfloat16,
+        indexes_kv_by_block_stride=True,
+    )
+    # The non-divisible draft layer needs padding but its backend does not
+    # support the strided padded-page view -> must raise, not silently pad.
+    draft_spec = new_sliding_window_spec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+        sliding_window=1024,
+        indexes_kv_by_block_stride=False,
+    )
+    specs = {"target_attn": target_spec, "draft_attn": draft_spec}
+
+    with pytest.raises(NotImplementedError):
+        kv_cache_utils.unify_kv_cache_spec_page_size(specs)
+
+
 def test_unify_hybrid_kv_cache_specs():
    # 1. has_full_attention and has_sliding_window
    before_spec_1 = new_kv_cache_spec()
@@ -144,6 +144,43 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
    assert req.request_id in output.num_scheduled_tokens


+def test_cached_request_data_resumed_all_token_ids_mrv1_only():
+    """all_token_ids carries a resumed request's token ids to the connector
+    for the V1 model runner, but is skipped entirely for the V2 model runner.
+    """
+    from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+
+    scheduler = create_scheduler()
+    (req,) = create_requests(num_requests=1, num_tokens=8)
+    req.append_output_token_ids([101, 102, 103])
+
+    # A resumed request was not scheduled in the previous step.
+    assert req.request_id not in scheduler.prev_step_scheduled_req_ids
+
+    empty_blocks = KVCacheBlocks(blocks=((),))
+
+    def make_cached():
+        return scheduler._make_cached_request_data(
+            running_reqs=[],
+            resumed_reqs=[req],
+            num_scheduled_tokens={req.request_id: 1},
+            spec_decode_tokens={},
+            req_to_new_blocks={req.request_id: empty_blocks},
+        )
+
+    # V1 model runner: the full token id list is propagated.
+    assert not scheduler.use_v2_model_runner
+    cached = make_cached()
+    assert req.request_id in cached.resumed_req_ids
+    assert cached.all_token_ids[req.request_id] == list(req.all_token_ids)
+
+    # V2 model runner: all_token_ids is skipped entirely.
+    scheduler.use_v2_model_runner = True
+    cached = make_cached()
+    assert req.request_id in cached.resumed_req_ids
+    assert cached.all_token_ids == {}
+
+
 def test_schedule_partial_requests():
    """Test scheduling behavior with partial requests.

@@ -4,9 +4,16 @@
 import pytest

 from vllm import LLM, SamplingParams
+from vllm.platforms import current_platform

 from ....utils import create_new_process_for_each_test

+if current_platform.is_rocm():
+    pytest.skip(
+        "Cascade attention backends FLASH_ATTN and FLASHINFER are notsupported on ROCm",
+        allow_module_level=True,
+    )
+

@create_new_process_for_each_test()
@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "FLASHINFER"])
@@ -425,7 +425,7 @@ def _run_eagle_correctness(
            if "deepseek" in model_setup[1].lower():
                m.setenv("VLLM_ROCM_USE_AITER", "1")
                m.delenv("VLLM_MLA_DISABLE", raising=False)
-                attention_config = {"backend": "TRITON_MLA"}
+                attention_config = {"backend": "ROCM_AITER_MLA"}
            else:
                m.setenv("VLLM_ROCM_USE_AITER", "1")

@@ -22,7 +22,7 @@ from vllm.v1.kv_offload.base import (
    OffloadingGaugeMetadata,
    OffloadingHistogramMetadata,
 )
-from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
+from vllm.v1.kv_offload.factory import OffloadingSpecFactory

 LOAD_BYTES = _TransferMetricName.LOAD_BYTES
 LOAD_TIME = _TransferMetricName.LOAD_TIME
@@ -33,6 +33,8 @@ STORE_SIZE = _TransferMetricName.STORE_SIZE
 STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
 PENDING_STORES = "vllm:kv_offload_pending_stores"
 LOOKUP_LATENCY = "vllm:kv_offload_lookup_latency_seconds"
+MY_COUNTER = "my_counter"
+MY_LABEL = "my_label"


 class _FakeMetric:
@@ -67,6 +69,20 @@ class _FakeVllmConfig:
        )


+def _spec_cls_with_metric_definitions(
+    metric_definitions: dict[str, Any],
+) -> type:
+    """Build a fake offloading spec class reporting the given metric
+    definitions, so tests don't need to patch the real CPU spec."""
+
+    class _FakeOffloadingSpec:
+        @staticmethod
+        def build_metric_definitions(extra_config):
+            return metric_definitions
+
+    return _FakeOffloadingSpec
+
+
 def _metric_metadata():
    return {
        LOAD_BYTES: OffloadingCounterMetadata(
@@ -96,9 +112,17 @@ def _metric_metadata():
        LOOKUP_LATENCY: OffloadingHistogramMetadata(
            documentation="lookup latency",
        ),
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
    }


+def _unlabeled(values: dict[str, Any], metric_name: str) -> Any:
+    return values[metric_name][()]
+
+
 def test_build_kv_connector_stats_with_none():
    """Test that build_kv_connector_stats returns empty stats when given None."""
    stats = OffloadingConnector.build_kv_connector_stats(data=None)
@@ -131,13 +155,13 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():
            STORES_SKIPPED: _MetricType.COUNTER,
        },
        _StatsKey.DATA: {
-            LOAD_BYTES: 24,
-            LOAD_TIME: 1.5,
-            LOAD_SIZE: [16, 8],
-            STORE_BYTES: 3,
-            STORE_TIME: 0.3,
-            STORE_SIZE: [1, 2],
-            STORES_SKIPPED: 5,
+            LOAD_BYTES: {(): 24},
+            LOAD_TIME: {(): 1.5},
+            LOAD_SIZE: {(): [16, 8]},
+            STORE_BYTES: {(): 3},
+            STORE_TIME: {(): 0.3},
+            STORE_SIZE: {(): [1, 2]},
+            STORES_SKIPPED: {(): 5},
        },
    }

@@ -145,22 +169,28 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():

    assert isinstance(stats, OffloadingConnectorStats)
    values = stats.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 24
-    assert values[LOAD_TIME] == 1.5
-    assert values[LOAD_SIZE] == [16, 8]
-    assert values[STORE_BYTES] == 3
-    assert values[STORE_TIME] == 0.3
-    assert values[STORE_SIZE] == [1, 2]
-    assert values[STORES_SKIPPED] == 5
+    assert _unlabeled(values, LOAD_BYTES) == 24
+    assert _unlabeled(values, LOAD_TIME) == 1.5
+    assert _unlabeled(values, LOAD_SIZE) == [16, 8]
+    assert _unlabeled(values, STORE_BYTES) == 3
+    assert _unlabeled(values, STORE_TIME) == 0.3
+    assert _unlabeled(values, STORE_SIZE) == [1, 2]
+    assert _unlabeled(values, STORES_SKIPPED) == 5


 def _make_stats_data(
    metric_data: dict[str, Any],
    metric_metadata: dict[str, Any],
 ) -> dict[str, Any]:
-    """Build a structured data dict from flat metric data and metadata."""
+    """Build a structured data dict from flat metric data and metadata.
+
+    Values for unlabeled metrics may be passed flat (wrapped here under the
+    empty label tuple); values for labeled metrics must already be passed as
+    a ``{labelvalues: value}`` map.
+    """
    metric_types = {}
-    for key in metric_data:
+    data = {}
+    for key, value in metric_data.items():
        md = metric_metadata[key]
        if isinstance(md, OffloadingCounterMetadata):
            metric_types[key] = _MetricType.COUNTER
@@ -168,9 +198,10 @@ def _make_stats_data(
            metric_types[key] = _MetricType.GAUGE
        elif isinstance(md, OffloadingHistogramMetadata):
            metric_types[key] = _MetricType.HISTOGRAM
+        data[key] = value if md.labelnames else {(): value}
    return {
        _StatsKey.TYPES: metric_types,
-        _StatsKey.DATA: metric_data,
+        _StatsKey.DATA: data,
    }


@@ -215,34 +246,106 @@ def test_aggregate_same_connector():

    assert result is stats1  # Should return self
    values = result.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 34
-    assert values[LOAD_TIME] == 2.6
-    assert values[LOAD_SIZE] == [16, 8, 3, 7]
-    assert values[STORE_BYTES] == 19
-    assert values[STORE_TIME] == 2.3
-    assert values[STORE_SIZE] == [1, 2, 16]
-    assert values[STORES_SKIPPED] == 4
-    assert values[PENDING_STORES] == 1
-    assert values[LOOKUP_LATENCY] == [0.1, 0.2, 0.3]
+    assert _unlabeled(values, LOAD_BYTES) == 34
+    assert _unlabeled(values, LOAD_TIME) == 2.6
+    assert _unlabeled(values, LOAD_SIZE) == [16, 8, 3, 7]
+    assert _unlabeled(values, STORE_BYTES) == 19
+    assert _unlabeled(values, STORE_TIME) == 2.3
+    assert _unlabeled(values, STORE_SIZE) == [1, 2, 16]
+    assert _unlabeled(values, STORES_SKIPPED) == 4
+    assert _unlabeled(values, PENDING_STORES) == 1
+    assert _unlabeled(values, LOOKUP_LATENCY) == [0.1, 0.2, 0.3]
+
+
+def test_aggregate_labeled_metrics():
+    metadata = _metric_metadata()
+    stats1 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 10,
+                    ("b",): 3,
+                },
+            },
+            metadata,
+        ),
+    )
+    stats2 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 7,
+                    ("c",): 5,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    stats1.aggregate(stats2)
+
+    values = stats1.data[_StatsKey.DATA][MY_COUNTER]
+    assert values[("a",)] == 17
+    assert values[("b",)] == 3
+    assert values[("c",)] == 5
+
+
+def test_aggregate_labeled_metric_missing_from_self():
+    """Aggregating a labeled metric that self doesn't have at all yet."""
+    metadata = _metric_metadata()
+    stats1 = OffloadingConnectorStats()
+    stats2 = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 7,
+                    ("b",): 5,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    stats1.aggregate(stats2)
+
+    values = stats1.data[_StatsKey.DATA][MY_COUNTER]
+    assert values[("a",)] == 7
+    assert values[("b",)] == 5
+    assert stats1.data[_StatsKey.TYPES][MY_COUNTER] == _MetricType.COUNTER
+
+
+def test_helper_methods_accept_labeled_metrics():
+    stats = OffloadingConnectorStats()
+
+    stats.increase_counter(MY_COUNTER, 3, ("a",))
+    stats.increase_counter(MY_COUNTER, 4, ("a",))
+    stats.set_gauge(PENDING_STORES, 2, ("b",))
+    stats.observe_histogram(LOOKUP_LATENCY, 0.1, ("b",))
+    stats.observe_histogram(LOOKUP_LATENCY, 0.2, ("b",))
+
+    values = stats.data[_StatsKey.DATA]
+    assert values[MY_COUNTER][("a",)] == 7
+    assert values[PENDING_STORES][("b",)] == 2
+    assert values[LOOKUP_LATENCY][("b",)] == [0.1, 0.2]


 def test_aggregate_merges_types():
    stats1 = OffloadingConnectorStats(
        data={
            _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
-            _StatsKey.DATA: {LOAD_BYTES: 1},
+            _StatsKey.DATA: {LOAD_BYTES: {(): 1}},
        },
    )
    stats2 = OffloadingConnectorStats(
        data={
            _StatsKey.TYPES: {PENDING_STORES: _MetricType.GAUGE},
-            _StatsKey.DATA: {PENDING_STORES: 2},
+            _StatsKey.DATA: {PENDING_STORES: {(): 2}},
        },
    )

    result = stats1.aggregate(stats2)

-    assert result.data[_StatsKey.DATA][PENDING_STORES] == 2
+    assert _unlabeled(result.data[_StatsKey.DATA], PENDING_STORES) == 2
    assert result.data[_StatsKey.TYPES][PENDING_STORES] == _MetricType.GAUGE


@@ -283,6 +386,26 @@ def test_reduce():
    assert reduced[f"{LOOKUP_LATENCY}_sum"] == sum([0.1, 0.2, 0.3])


+def test_reduce_labeled_metrics():
+    metadata = _metric_metadata()
+    stats = OffloadingConnectorStats(
+        data=_make_stats_data(
+            {
+                MY_COUNTER: {
+                    ("a",): 17,
+                    ("b",): 3,
+                },
+            },
+            metadata,
+        ),
+    )
+
+    reduced = stats.reduce()
+
+    assert reduced[f"{MY_COUNTER}:{('a',)}"] == 17
+    assert reduced[f"{MY_COUNTER}:{('b',)}"] == 3
+
+
 def test_reset():
    """Test that reset() resets all connector stats."""
    metadata = _metric_metadata()
@@ -326,11 +449,11 @@ def test_prom_metrics_observes_manager_counter():
    prom_metrics.observe(
        {
            _StatsKey.TYPES: {STORES_SKIPPED: _MetricType.COUNTER},
-            _StatsKey.DATA: {STORES_SKIPPED: 7},
+            _StatsKey.DATA: {STORES_SKIPPED: {(): 7}},
        }
    )

-    counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED)]
+    counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED, ())]
    assert counter.increments == [7]
    counter_def = prom_metrics._offloading_metric_defs[STORES_SKIPPED]
    assert counter_def.kwargs["name"] == "vllm:kv_offload_stores_skipped"
@@ -360,22 +483,22 @@ def test_prom_metrics_observes_flat_transfer_metrics_and_legacy_metrics():
                STORE_SIZE: _MetricType.HISTOGRAM,
            },
            _StatsKey.DATA: {
-                LOAD_BYTES: 24,
-                LOAD_TIME: 1.5,
-                LOAD_SIZE: [16, 8],
-                STORE_BYTES: 3,
-                STORE_TIME: 0.3,
-                STORE_SIZE: [1, 2],
+                LOAD_BYTES: {(): 24},
+                LOAD_TIME: {(): 1.5},
+                LOAD_SIZE: {(): [16, 8]},
+                STORE_BYTES: {(): 3},
+                STORE_TIME: {(): 0.3},
+                STORE_SIZE: {(): [1, 2]},
            },
        }
    )

-    assert prom_metrics.offloading_metrics[(0, LOAD_BYTES)].increments == [24]
-    assert prom_metrics.offloading_metrics[(0, LOAD_TIME)].increments == [1.5]
-    assert prom_metrics.offloading_metrics[(0, LOAD_SIZE)].observed == [16, 8]
-    assert prom_metrics.offloading_metrics[(0, STORE_BYTES)].increments == [3]
-    assert prom_metrics.offloading_metrics[(0, STORE_TIME)].increments == [0.3]
-    assert prom_metrics.offloading_metrics[(0, STORE_SIZE)].observed == [1, 2]
+    assert prom_metrics.offloading_metrics[(0, LOAD_BYTES, ())].increments == [24]
+    assert prom_metrics.offloading_metrics[(0, LOAD_TIME, ())].increments == [1.5]
+    assert prom_metrics.offloading_metrics[(0, LOAD_SIZE, ())].observed == [16, 8]
+    assert prom_metrics.offloading_metrics[(0, STORE_BYTES, ())].increments == [3]
+    assert prom_metrics.offloading_metrics[(0, STORE_TIME, ())].increments == [0.3]
+    assert prom_metrics.offloading_metrics[(0, STORE_SIZE, ())].observed == [1, 2]

    assert prom_metrics.counter_kv_bytes[(0, "CPU_to_GPU")].increments == [24]
    assert prom_metrics.counter_kv_transfer_time[(0, "CPU_to_GPU")].increments == [1.5]
@@ -396,7 +519,9 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
        ),
    }
    with patch.object(
-        CPUOffloadingSpec, "build_metric_definitions", return_value=metric_definitions
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
    ):
        prom_metrics = OffloadPromMetrics(
            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
@@ -416,20 +541,91 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
                LOOKUP_LATENCY: _MetricType.HISTOGRAM,
            },
            _StatsKey.DATA: {
-                PENDING_STORES: 5,
-                LOOKUP_LATENCY: [0.2, 0.4],
+                PENDING_STORES: {(): 5},
+                LOOKUP_LATENCY: {(): [0.2, 0.4]},
            },
        }
    )

-    gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES)]
-    histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY)]
+    gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES, ())]
+    histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY, ())]
    assert gauge.set_values == [5]
    assert histogram.observed == [0.2, 0.4]
    histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY]
    assert histogram_def.kwargs["buckets"] == (0.1, 1.0)


+def test_prom_metrics_lazily_observes_labeled_metric():
+    metric_definitions = {
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
+    }
+    with patch.object(
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
+    ):
+        prom_metrics = OffloadPromMetrics(
+            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
+            metric_types={
+                Gauge: _FakeMetric,
+                Counter: _FakeMetric,
+                Histogram: _FakeMetric,
+            },
+            labelnames=["model_name", "engine"],
+            per_engine_labelvalues={0: ["model", "0"]},
+        )
+
+    assert (0, MY_COUNTER, ("a",)) not in prom_metrics.offloading_metrics
+
+    prom_metrics.observe(
+        {
+            _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
+            _StatsKey.DATA: {MY_COUNTER: {("a",): 7}},
+        }
+    )
+
+    counter = prom_metrics.offloading_metrics[(0, MY_COUNTER, ("a",))]
+    assert counter.increments == [7]
+    assert counter.labelvalues == ("model", "0", "a")
+    counter_def = prom_metrics._offloading_metric_defs[MY_COUNTER]
+    assert counter_def.kwargs["labelnames"] == ["model_name", "engine", MY_LABEL]
+
+
+def test_prom_metrics_rejects_wrong_label_count():
+    metric_definitions = {
+        MY_COUNTER: OffloadingCounterMetadata(
+            documentation="counter with a label",
+            labelnames=(MY_LABEL,),
+        ),
+    }
+    with patch.object(
+        OffloadingSpecFactory,
+        "get_spec_cls",
+        return_value=_spec_cls_with_metric_definitions(metric_definitions),
+    ):
+        prom_metrics = OffloadPromMetrics(
+            vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
+            metric_types={
+                Gauge: _FakeMetric,
+                Counter: _FakeMetric,
+                Histogram: _FakeMetric,
+            },
+            labelnames=["model_name", "engine"],
+            per_engine_labelvalues={0: ["model", "0"]},
+        )
+
+    with pytest.raises(AssertionError, match="expects 1 labels"):
+        prom_metrics.observe(
+            {
+                _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
+                _StatsKey.DATA: {MY_COUNTER: {("a", "extra"): 7}},
+            }
+        )
+
+
 def test_prom_metrics_uses_configured_manager_metrics():
    prom_metrics = OffloadPromMetrics(
        vllm_config=_FakeVllmConfig(store_threshold=0),  # type: ignore[arg-type]
@@ -458,9 +654,9 @@ def test_aggregate_into_empty_stats():
                PENDING_STORES: _MetricType.GAUGE,
            },
            _StatsKey.DATA: {
-                LOAD_BYTES: 42,
-                LOAD_SIZE: [10, 20],
-                PENDING_STORES: 3,
+                LOAD_BYTES: {(): 42},
+                LOAD_SIZE: {(): [10, 20]},
+                PENDING_STORES: {(): 3},
            },
        },
    )
@@ -469,9 +665,9 @@ def test_aggregate_into_empty_stats():

    assert result is empty
    values = result.data[_StatsKey.DATA]
-    assert values[LOAD_BYTES] == 42
-    assert values[LOAD_SIZE] == [10, 20]
-    assert values[PENDING_STORES] == 3
+    assert _unlabeled(values, LOAD_BYTES) == 42
+    assert _unlabeled(values, LOAD_SIZE) == [10, 20]
+    assert _unlabeled(values, PENDING_STORES) == 3


 def test_prom_metrics_multi_engine_routing():
@@ -490,14 +686,13 @@ def test_prom_metrics_multi_engine_routing():
    prom_metrics.observe(
        {
            _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
-            _StatsKey.DATA: {LOAD_BYTES: 100},
+            _StatsKey.DATA: {LOAD_BYTES: {(): 100}},
        },
        engine_idx=1,
    )

-    engine0 = prom_metrics.offloading_metrics[(0, LOAD_BYTES)]
-    engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES)]
-    assert engine0.increments == []
+    assert (0, LOAD_BYTES, ()) not in prom_metrics.offloading_metrics
+    engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES, ())]
    assert engine1.increments == [100]


@@ -518,6 +713,6 @@ def test_prom_metrics_rejects_undeclared_metric():
        prom_metrics.observe(
            {
                _StatsKey.TYPES: {"unknown:metric": _MetricType.COUNTER},
-                _StatsKey.DATA: {"unknown:metric": 1},
+                _StatsKey.DATA: {"unknown:metric": {(): 1}},
            }
        )
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Iterable
+from types import SimpleNamespace
 from unittest.mock import MagicMock

 import pytest
@@ -1278,11 +1279,11 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
    )

    finalized: list[str] = []
-    runner.manager.on_request_finished.side_effect = (
-        lambda req_context: finalized.append(req_context.req_id)
+    runner.manager.on_request_finished.side_effect = lambda req_context: (
+        finalized.append(req_context.req_id)
    )
-    runner.manager.prepare_store.side_effect = (
-        lambda keys, req_context: generate_store_output(keys)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
    )

    # Decode a couple of blocks and keep every transfer in flight, so the
@@ -1314,6 +1315,100 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
    assert req_id not in cs._req_status


+def test_pending_transfer_defers_prefix_lookup():
+    """A request with an in-flight store must not issue a load on re-admission.
+
+    With async scheduling, a preempted request's store can be flushed by the
+    worker before the scheduler consumes its completion. If the request is
+    re-admitted in that window, the connector should defer it instead of
+    looking up offloaded blocks and later asserting when a load is queued while
+    the store job is still tracked.
+    """
+    scheduler = object.__new__(OffloadingConnectorScheduler)
+    scheduler.manager = MagicMock(spec=OffloadingManager)
+
+    request = SimpleNamespace(request_id="req-0")
+    group_state = SimpleNamespace(block_ids=[1, 2, 3])
+    req_status = SimpleNamespace(
+        group_states=[group_state],
+        transfer_jobs={123},
+    )
+    scheduler._req_status = {request.request_id: req_status}
+
+    matched_tokens, is_async = scheduler.get_num_new_matched_tokens(
+        request,
+        num_computed_tokens=0,
+    )
+
+    assert matched_tokens is None
+    assert is_async is False
+    assert group_state.block_ids == []
+    scheduler.manager.lookup.assert_not_called()
+
+
+def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner):
+    """A preempted request can be scheduled again before flush output is read.
+
+    EngineCore.step_with_batch_queue() may schedule a new batch while a prior
+    preemption batch is still queued. The store completion from jobs_to_flush is
+    only cleared when that queued output reaches update_from_output(), so the
+    re-admission path must defer while the scheduler still tracks the store.
+    """
+    block_size = 4
+    block_size_factor = 3
+    offloaded_block_size = block_size * block_size_factor
+
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=100,
+        async_scheduling=True,
+        block_size_factor=block_size_factor,
+    )
+    free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue
+    num_free_blocks_empty = free_block_queue.num_free_blocks
+
+    req_id = "0"
+    runner.new_request(token_ids=[0] * offloaded_block_size * 2)
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    runner.run(decoded_tokens=[0], complete_transfers=False)
+    runner.run(
+        decoded_tokens=[0] * (2 * offloaded_block_size - block_size),
+        complete_transfers=False,
+    )
+
+    req_status = runner.connector_scheduler._req_status[req_id]
+    pending_store_jobs = set(req_status.transfer_jobs)
+    assert pending_store_jobs
+    assert all(
+        runner.connector_scheduler._jobs[jid].is_store for jid in pending_store_jobs
+    )
+
+    free_block_queue.num_free_blocks = 0
+    preempt_output = runner.scheduler.schedule()
+    assert preempt_output.preempted_req_ids == {req_id}
+    assert preempt_output.kv_connector_metadata is not None
+    assert pending_store_jobs <= preempt_output.kv_connector_metadata.jobs_to_flush
+    assert req_status.transfer_jobs == pending_store_jobs
+
+    # Simulate the async batch-queue window: schedule again before the
+    # preemption batch's ModelRunnerOutput is consumed by update_from_output().
+    free_block_queue.num_free_blocks = num_free_blocks_empty
+    assert runner.scheduler.reset_prefix_cache()
+    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: len(
+        key
+    )
+
+    readmit_output = runner.scheduler.schedule()
+
+    assert readmit_output.num_scheduled_tokens == {}
+    assert readmit_output.kv_connector_metadata is not None
+    assert readmit_output.kv_connector_metadata.load_jobs == {}
+    assert req_status.transfer_jobs == pending_store_jobs
+
+
@pytest.mark.parametrize("async_scheduling", [True, False])
 def test_swa_alignment_skip(request_runner, async_scheduling: bool):
    """SWA blocks unreachable by the load path are skipped during store.
@@ -7,7 +7,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp
    ExternalCachedBlockPool,
    MooncakeStoreCoordinator,
 )
-from vllm.v1.core.kv_cache_utils import BlockHash, BlockHashListWithBlockSize
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+    chunk_hashes_for_block_size,
+)
+from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheGroupSpec,
@@ -182,7 +185,7 @@ def test_coordinator_group_block_size_double_hash():
    ]
    coord = _make_coord(groups, hash_block_size=16)
    hs = _hashes(4)
-    big_hashes = list(BlockHashListWithBlockSize(hs, 16, 32))
+    big_hashes = list(chunk_hashes_for_block_size(hs, 16, 32))
    exists = {(0, bytes(h)) for h in hs}
    exists |= {(1, bytes(bh)) for bh in big_hashes}
    cmap = ExternalCachedBlockPool(exists)
@@ -323,8 +323,8 @@ def test_recv_skips_swa_blocks_before_window():

 def test_chunked_token_database_hash_block_size_smaller_than_block_size():
    """DSv4-style: hash_block_size=4, group block_size=16 — process_tokens
-    must merge every 4 fine hashes into one chunk hash via
-    BlockHashListWithBlockSize."""
+    keys each 16-token chunk by its last fine hash, keeping the Mooncake key
+    at one digest instead of concatenating all 4 fine hashes."""
    md = KeyMetadata("m", 0, 0, 0, 0, group_id=3)
    db = ChunkedTokenDatabase(md, block_size=16, hash_block_size=4)
    db.set_kv_caches_base_addr([0])
@@ -335,8 +335,7 @@ def test_chunked_token_database_hash_block_size_smaller_than_block_size():
    assert len(out) == 2
    assert out[0][0] == 0 and out[0][1] == 16
    assert out[1][0] == 16 and out[1][1] == 32
-    # Each chunk's hash is the concatenation of 4 fine hashes.
-    expected0 = b"".join(fine_hashes[0:4]).hex()
-    expected1 = b"".join(fine_hashes[4:8]).hex()
-    assert out[0][2].chunk_hash == expected0
-    assert out[1][2].chunk_hash == expected1
+    # Each chunk's hash is its last (4th) fine hash, which already chains the
+    # prior three.
+    assert out[0][2].chunk_hash == fine_hashes[3].hex()
+    assert out[1][2].chunk_hash == fine_hashes[7].hex()
@@ -23,6 +23,7 @@ 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 (
+    BlobBlockHashes,
    ChunkedTokenDatabase,
    KeyMetadata,
    LoadSpec,
@@ -32,6 +33,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
    MooncakeStoreConnectorStats,
 )
+from vllm.v1.core.kv_cache_utils import BlockHash


 def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
@@ -1179,9 +1181,9 @@ def test_store_sending_thread_kv_events_use_group_chunk_metadata():
    assert full_event.group_idx == 0
    assert full_event.block_size == 32
    assert full_event.token_ids == list(range(32))
-    assert full_event.block_hashes == [
-        maybe_convert_block_hash(BlockHash(b"".join(hs)))
-    ]
+    # block_size=32 over hash_block_size=8 (scale 4): the chunk is keyed by its
+    # last sub-hash, not the concatenation of all four.
+    assert full_event.block_hashes == [maybe_convert_block_hash(BlockHash(hs[3]))]

    assert swa_event.group_idx == 1
    assert swa_event.block_size == 8
@@ -1749,3 +1751,33 @@ def test_store_worker_close_swallows_store_errors():
    worker.close()

    assert worker.store is None
+
+
+def test_blob_block_hashes_wire_roundtrip():
+    """The lookup wire format sends a ``hash_len`` frame plus the raw hashes
+    concatenated back-to-back; the server rebuilds them through a zero-copy
+    ``BlobBlockHashes`` view over the frame buffer."""
+    hashes = [BlockHash(bytes([i]) * 16) for i in range(5)]
+    hash_len = len(hashes[0])
+
+    # Client side (LookupKeyClient._lookup): flat payload frame.
+    blob = b"".join(hashes)
+
+    # Server side (LookupKeyServer): view over the frame buffer (a memoryview),
+    # never materializing the full hash list upfront.
+    view = BlobBlockHashes(memoryview(blob), hash_len)
+
+    assert len(view) == 5
+    assert list(view) == hashes  # default Sequence iter terminates via IndexError
+    assert [bytes(h) for h in view] == hashes
+    assert bytes(view[-1]) == hashes[-1]
+    assert [bytes(h) for h in view[1:3]] == hashes[1:3]
+    with pytest.raises(IndexError):
+        _ = view[5]
+
+
+def test_blob_block_hashes_empty():
+    """Empty lookups send hash_len=0 and an empty payload."""
+    view = BlobBlockHashes(memoryview(b""), 0)
+    assert len(view) == 0
+    assert list(view) == []
@@ -36,13 +36,33 @@ from vllm.utils.network_utils import (
    get_ip,
    make_zmq_path,
 )
-from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+)

 from .utils import create_request, create_scheduler


 def _make_test_kv_cache_config() -> KVCacheConfig:
-    return KVCacheConfig(num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[])
+    layer_names = ["layer0", "layer1", "layer2"]
+    return KVCacheConfig(
+        num_blocks=2,
+        kv_cache_tensors=[KVCacheTensor(size=0, shared_by=layer_names)],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                layer_names=layer_names,
+                kv_cache_spec=FullAttentionSpec(
+                    block_size=16,
+                    num_kv_heads=4,
+                    head_size=64,
+                    dtype=torch.float16,
+                ),
+            )
+        ],
+    )


 aiter_available = importlib.util.find_spec("aiter") is not None
@@ -175,9 +195,18 @@ class FakeMoRIIOConnectorWorker(MoRIIOConnectorWorker):
    REMOTE_ENGINE_ID = "remote_engine"

    def __init__(
-        self, *args, hand_shake_latency: float = 1.8, kv_cache_layout="HND", **kwargs
+        self,
+        vllm_config,
+        engine_id,
+        *args,
+        hand_shake_latency: float = 1.8,
+        kv_cache_layout="HND",
+        kv_cache_config=None,
+        **kwargs,
    ):
-        super().__init__(*args, **kwargs)
+        super().__init__(
+            vllm_config, engine_id, kv_cache_config or _make_test_kv_cache_config()
+        )


 def create_vllm_config(
@@ -0,0 +1,228 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import importlib.util
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
+
+aiter_available = importlib.util.find_spec("aiter") is not None
+mori_available = importlib.util.find_spec("mori") is not None
+
+if not (current_platform.is_rocm() and mori_available):
+    pytest.skip(
+        "MoRIIOs are only available on ROCm with mori package installed",
+        allow_module_level=True,
+    )
+
+moriio_layout = importlib.import_module(
+    "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout"
+)
+
+
+def _full_spec(block_size: int = 4) -> FullAttentionSpec:
+    return FullAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=2,
+        head_size=3,
+        dtype=torch.bfloat16,
+    )
+
+
+def _mla_spec(block_size: int = 4) -> MLAAttentionSpec:
+    return MLAAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=3,
+        dtype=torch.bfloat16,
+    )
+
+
+def _worker(
+    kv_caches: dict[str, torch.Tensor],
+    layer_to_spec: dict[str, object],
+    num_blocks: int = 8,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        kv_caches=kv_caches,
+        layer_to_spec=layer_to_spec,
+        num_blocks=num_blocks,
+        block_size=4,
+    )
+
+
+def _remote_meta(num_blocks: int = 16) -> SimpleNamespace:
+    return SimpleNamespace(num_blocks=num_blocks)
+
+
+def test_separated_kv_layout_uses_kv_axis_zero_and_block_axis_one():
+    cache = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    geometry = moriio_layout.get_layer_transfer_geometry(
+        "layer", cache, worker.layer_to_spec, remote_num_blocks=16
+    )
+    assert geometry.block_stride == 24
+    assert geometry.local_kv_stride == 192
+    assert geometry.remote_kv_stride == 384
+    assert geometry.split_kv_regions
+
+    assert moriio_layout.compute_block_transfer_offsets(
+        "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
+    ) == ([48, 144, 432, 528], [192, 240, 960, 1008], [48, 48, 48, 48])
+
+
+def test_interleaved_kv_layout_uses_block_axis_zero_and_kv_axis_one():
+    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    geometry = moriio_layout.get_layer_transfer_geometry(
+        "layer", cache, worker.layer_to_spec, remote_num_blocks=16
+    )
+    assert geometry.block_stride == 48
+    assert geometry.local_kv_stride == 24
+    assert geometry.remote_kv_stride == 24
+    assert not geometry.split_kv_regions
+
+    assert moriio_layout.compute_block_transfer_offsets(
+        "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
+    ) == ([96, 288], [384, 480], [96, 96])
+
+
+def test_mla_key_only_layout_transfers_one_slab_per_block():
+    cache = torch.empty((8, 4, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _mla_spec()})
+
+    geometry = moriio_layout.get_layer_transfer_geometry(
+        "layer", cache, worker.layer_to_spec, remote_num_blocks=16
+    )
+    assert geometry.block_stride == 12
+    assert geometry.local_kv_stride is None
+    assert geometry.remote_kv_stride is None
+    assert geometry.transfers_per_block == 1
+
+    assert moriio_layout.compute_block_transfer_offsets(
+        "layer", cache, worker.layer_to_spec, [1, 3], [4, 5], _remote_meta().num_blocks
+    ) == ([24, 72], [96, 120], [24, 24])
+
+
+def test_mixed_layers_compute_distinct_offsets_per_layer():
+    kv_caches = {
+        "separated": torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16),
+        "interleaved": torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16),
+        "indexer": torch.empty((8, 4, 3), dtype=torch.bfloat16),
+    }
+    worker = _worker(
+        kv_caches,
+        {
+            "separated": _full_spec(),
+            "interleaved": _full_spec(),
+            "indexer": _mla_spec(),
+        },
+    )
+
+    separated = moriio_layout.compute_block_transfer_offsets(
+        "separated",
+        kv_caches["separated"],
+        worker.layer_to_spec,
+        [1, 3],
+        [4, 5],
+        _remote_meta().num_blocks,
+    )
+    interleaved = moriio_layout.compute_block_transfer_offsets(
+        "interleaved",
+        kv_caches["interleaved"],
+        worker.layer_to_spec,
+        [1, 3],
+        [4, 5],
+        _remote_meta().num_blocks,
+    )
+    indexer = moriio_layout.compute_block_transfer_offsets(
+        "indexer",
+        kv_caches["indexer"],
+        worker.layer_to_spec,
+        [1, 3],
+        [4, 5],
+        _remote_meta().num_blocks,
+    )
+
+    assert separated != interleaved
+    assert separated != indexer
+    assert interleaved != indexer
+
+
+def test_block_id_length_mismatch_raises_value_error():
+    cache = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    with pytest.raises(ValueError, match="must have the same length"):
+        moriio_layout.compute_block_transfer_offsets(
+            "layer", cache, worker.layer_to_spec, [1, 3], [4], _remote_meta().num_blocks
+        )
+
+
+def test_registration_regions_do_not_split_interleaved_or_mla_cache():
+    separated = torch.empty((2, 8, 4, 2, 3), dtype=torch.bfloat16)
+    interleaved = torch.empty((8, 2, 4, 2, 3), dtype=torch.bfloat16)
+    indexer = torch.empty((8, 4, 3), dtype=torch.bfloat16)
+    worker = _worker(
+        {
+            "separated": separated,
+            "interleaved": interleaved,
+            "indexer": indexer,
+        },
+        {
+            "separated": _full_spec(),
+            "interleaved": _full_spec(),
+            "indexer": _mla_spec(),
+        },
+    )
+
+    separated_regions = moriio_layout.iter_layer_registration_regions(
+        "separated", separated, worker.layer_to_spec
+    )
+    interleaved_regions = moriio_layout.iter_layer_registration_regions(
+        "interleaved", interleaved, worker.layer_to_spec
+    )
+    indexer_regions = moriio_layout.iter_layer_registration_regions(
+        "indexer", indexer, worker.layer_to_spec
+    )
+
+    assert [region[0].data_ptr() for region in separated_regions] == [
+        separated[0].data_ptr(),
+        separated[1].data_ptr(),
+    ]
+    assert separated_regions[0][1] == 8 * 48
+    assert separated_regions[1][1] == 8 * 48
+
+    assert len(interleaved_regions) == 1
+    assert interleaved_regions[0][0].data_ptr() == interleaved.data_ptr()
+    assert interleaved_regions[0][1] == 8 * 2 * 48
+
+    assert len(indexer_regions) == 1
+    assert indexer_regions[0][0].data_ptr() == indexer.data_ptr()
+    assert indexer_regions[0][1] == 8 * 24
+
+
+def test_registration_regions_use_layer_num_blocks():
+    cache = torch.empty((4, 2, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()}, num_blocks=8)
+
+    regions = moriio_layout.iter_layer_registration_regions(
+        "layer", cache, worker.layer_to_spec
+    )
+
+    assert len(regions) == 1
+    assert regions[0][1] == 4 * 2 * 48
+
+
+def test_unsupported_shape_raises_value_error():
+    cache = torch.empty((8, 4, 2, 3), dtype=torch.bfloat16)
+    worker = _worker({"layer": cache}, {"layer": _full_spec()})
+
+    with pytest.raises(ValueError, match="Unsupported MoRIIO K/V cache shape"):
+        moriio_layout.get_layer_transfer_geometry("layer", cache, worker.layer_to_spec)
@@ -14,12 +14,13 @@ from vllm.v1.kv_offload.base import (
    ReqContext,
    make_offload_key,
 )
-from vllm.v1.kv_offload.cpu.common import CPULoadStoreSpec
+from vllm.v1.kv_offload.cpu.common import (
+    CPULoadStoreSpec,
+    CPUOffloadingMetrics,
+)
 from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
 from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy

-STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
-

 def make_req_context(
    req_id: str = "", kv_transfer_params: dict | None = None
@@ -181,10 +182,45 @@ def test_filter_reused_manager_reports_stores_skipped_counter():
    )
    stats = manager.get_stats()
    assert stats is not None
-    assert stats.reduce()[STORES_SKIPPED] == 3
+    assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 3
    stats = manager.get_stats()
    assert stats is not None
-    assert stats.reduce()[STORES_SKIPPED] == 0
+    assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 0
+
+
+def test_cpu_manager_reports_cache_usage_gauge():
+    def check_usage_stats(manager: CPUOffloadingManager, value: float):
+        stats = manager.get_stats()
+        assert stats is not None
+        assert stats.reduce()[
+            CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC
+        ] == pytest.approx(value)
+
+    # Zero-capacity manager always reports 0.0
+    manager = make_cpu_manager(num_blocks=0)
+    check_usage_stats(manager, 0.0)
+
+    # Empty manager (4 blocks, none allocated): usage = 0.0
+    manager = make_cpu_manager(num_blocks=4)
+    check_usage_stats(manager, 0.0)
+
+    # After allocating 2 of 4 blocks: usage = 0.5
+    manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 0.5)
+
+    # After filling all 4 blocks: usage = 1.0
+    manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 1.0)
+
+    # After completing store, the blocks becomes evictable as it is not actively used
+    # and usage drops.
+    manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 0.5)
+
+    # After completing store, the blocks becomes evictable as it is not actively used
+    # and usage drops.
+    manager.complete_store(to_keys([3, 4]), _EMPTY_REQ_CTX)
+    check_usage_stats(manager, 0.0)


 def test_cpu_manager():
@@ -145,7 +145,6 @@ def _generate_fake_sampling_metadata(
        vllm_config.scheduler_config.max_num_seqs,
        num_spec,
        device,
-        PIN_MEMORY_AVAILABLE,
    )
    fake_sampling_metadata = SamplingMetadata(
        temperature=torch.full((batch_size,), 0.0),
@@ -880,7 +879,6 @@ def test_maybe_create_thinking_budget_holder_without_reasoning():
            cfg.scheduler_config.max_num_seqs,
            0,
            torch.device("cpu"),
-            False,
        )
        is None
    )
@@ -6,6 +6,7 @@ from __future__ import annotations

 from dataclasses import dataclass

+import pytest
 import torch

 from vllm import SamplingParams
@@ -1528,3 +1529,232 @@ def test_reset_pending_loads() -> None:
    # All GPU blocks free
    num_used = gpu_pool.num_gpu_blocks - gpu_pool.get_num_free_blocks()
    assert num_used == 1, f"Expected only null block in use, got {num_used}"
+
+
+def _make_cp_vllm_config(
+    dcp_world_size: int = 1,
+    pcp_world_size: int = 1,
+) -> VllmConfig:
+    """VllmConfig with context-parallel sizes set for scheduler-only tests."""
+    cfg = _make_vllm_config()
+
+    cfg.parallel_config.decode_context_parallel_size = dcp_world_size
+    cfg.parallel_config.prefill_context_parallel_size = pcp_world_size
+    return cfg
+
+
+def _make_cp_scheduler(
+    *,
+    dcp_world_size: int = 1,
+    pcp_world_size: int = 1,
+    num_cpu_blocks: int = 8,
+    num_gpu_blocks: int = 16,
+    lazy: bool = False,
+) -> SchedulerFixture:
+    """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size."""
+    cp_world_size = dcp_world_size * pcp_world_size
+    virtual_block_size = BLOCK_SIZE * cp_world_size
+
+    kv_cache_config = _make_kv_cache_config(num_gpu_blocks)
+    vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size)
+    cpu_capacity_bytes = _BYTES_PER_BLOCK * num_cpu_blocks
+
+    sched = SimpleCPUOffloadScheduler(
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+        cpu_capacity_bytes=cpu_capacity_bytes,
+        scheduler_block_size=virtual_block_size,
+        hash_block_size=virtual_block_size,
+        lazy_offload=lazy,
+    )
+
+    gpu_block_pool = BlockPool(
+        num_gpu_blocks=num_gpu_blocks,
+        enable_caching=True,
+        hash_block_size=virtual_block_size,
+    )
+    sched.bind_gpu_block_pool(gpu_block_pool)
+
+    return SchedulerFixture(
+        scheduler=sched,
+        gpu_block_pool=gpu_block_pool,
+        vllm_config=vllm_config,
+        kv_cache_config=kv_cache_config,
+    )
+
+
+def _make_cp_request(
+    num_blocks: int,
+    virtual_block_size: int,
+    request_id: str | None = None,
+) -> Request:
+    """Create a request whose block hashes are computed at the virtual
+    (CP-scaled) block size, matching what the real scheduler does.
+    """
+    global _req_counter
+    _req_counter += 1
+    if request_id is None:
+        request_id = f"req-cp-{_req_counter}"
+
+    num_tokens = num_blocks * virtual_block_size + 1
+    start = _req_counter * 10000
+    prompt_token_ids = list(range(start, start + num_tokens))
+    sampling_params = SamplingParams(max_tokens=1)
+
+    return Request(
+        request_id=request_id,
+        prompt_token_ids=prompt_token_ids,
+        sampling_params=sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=get_request_block_hasher(virtual_block_size, sha256),
+    )
+
+
+def _allocate_cp_gpu_blocks(
+    gpu_block_pool: BlockPool,
+    request: Request,
+    num_blocks: int,
+    virtual_block_size: int,
+    group_id: int = 0,
+) -> list:
+    """Allocate GPU blocks and cache them using the CP-scaled block size."""
+    blocks = gpu_block_pool.get_new_blocks(num_blocks)
+    num_full = min(num_blocks, len(request.block_hashes))
+    if num_full > 0:
+        gpu_block_pool.cache_full_blocks(
+            request=request,
+            blocks=blocks,
+            num_cached_blocks=0,
+            num_full_blocks=num_full,
+            block_size=virtual_block_size,
+            kv_cache_group_id=group_id,
+        )
+    return blocks
+
+
+# ---------------------------------------------------------------------------
+# Test 15: CP block size scaling is correct
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "dcp_world_size, pcp_world_size",
+    [
+        (2, 1),  # DCP only
+        (1, 2),  # PCP only
+        (2, 2),  # DCP + PCP
+    ],
+)
+def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None:
+    """Verify that the scheduler's block_size and cp_world_size are correctly
+    scaled when context parallelism is enabled."""
+    fix = _make_cp_scheduler(
+        dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
+    )
+    sched = fix.scheduler
+
+    expected_cp = dcp_world_size * pcp_world_size
+    assert sched.cp_world_size == expected_cp
+    assert sched.block_size == BLOCK_SIZE * expected_cp
+
+
+# ---------------------------------------------------------------------------
+# Test 16: CP eager store-and-load roundtrip
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize(
+    "dcp_world_size, pcp_world_size",
+    [
+        (2, 1),
+        (1, 2),
+    ],
+)
+def test_cp_eager_store_and_load_roundtrip(
+    dcp_world_size: int, pcp_world_size: int
+) -> None:
+    """With CP enabled, store blocks to CPU and reload them for a new request
+    with matching tokens.  Verifies that hash matching and transfer-pair
+    construction work with the virtual block size."""
+    fix = _make_cp_scheduler(
+        dcp_world_size=dcp_world_size,
+        pcp_world_size=pcp_world_size,
+        num_cpu_blocks=8,
+        num_gpu_blocks=16,
+        lazy=False,
+    )
+    sched = fix.scheduler
+    cp = dcp_world_size * pcp_world_size
+    vbs = BLOCK_SIZE * cp
+
+    num_blocks = 2
+    req = _make_cp_request(num_blocks, vbs)
+
+    # Allocate GPU blocks and register hashes
+    gpu_blocks = _allocate_cp_gpu_blocks(fix.gpu_block_pool, req, num_blocks, vbs)
+    kv_blocks = KVCacheBlocks(blocks=(gpu_blocks,))
+    req.num_computed_tokens = num_blocks * vbs
+    sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
+
+    block_ids = kv_blocks.get_block_ids()
+    sched_out = make_scheduler_output(
+        {req.request_id: num_blocks * vbs},
+        new_reqs={req.request_id: block_ids},
+    )
+
+    meta = sched.build_connector_meta(sched_out)
+    assert meta.store_event >= 0, "Expected a store event"
+    assert len(meta.store_gpu_blocks) == num_blocks
+    assert len(meta.store_cpu_blocks) == num_blocks
+    simulate_store_completion(sched, meta.store_event)
+
+    # New request with same tokens — should get a full CPU cache hit.
+    req2 = Request(
+        request_id="req-cp-load",
+        prompt_token_ids=req.prompt_token_ids,
+        sampling_params=req.sampling_params,
+        pooling_params=None,
+        mm_features=None,
+        block_hasher=req._block_hasher,
+    )
+
+    hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
+    assert hit_tokens == num_blocks * vbs
+    assert is_async is True
+
+    # Allocate fresh GPU blocks for the load.
+    gpu_blocks2 = fix.gpu_block_pool.get_new_blocks(num_blocks)
+    kv_blocks2 = KVCacheBlocks(blocks=(gpu_blocks2,))
+    sched.update_state_after_alloc(req2, kv_blocks2, num_external_tokens=hit_tokens)
+
+    sched_out2 = make_scheduler_output(
+        {req2.request_id: 1},
+        new_reqs={req2.request_id: kv_blocks2.get_block_ids()},
+    )
+    meta2 = sched.build_connector_meta(sched_out2)
+    assert meta2.load_event >= 0, "Expected a load event"
+    assert len(meta2.load_gpu_blocks) == num_blocks
+    assert len(meta2.load_cpu_blocks) == num_blocks
+
+
+# ---------------------------------------------------------------------------
+# Test 17: CP lazy target blocks are scaled correctly
+# ---------------------------------------------------------------------------
+@pytest.mark.parametrize("cp_world_size", [1, 2, 4])
+def test_cp_lazy_target_blocks_scaling(cp_world_size: int) -> None:
+    """_estimate_lazy_target_blocks returns fewer blocks when cp_world_size > 1
+    because each virtual block covers more tokens."""
+    kv_cache_config = _make_kv_cache_config(num_blocks=16)
+    max_batched = 64
+
+    target_base = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        kv_cache_config, max_batched, cp_world_size=1
+    )
+    target_cp = SimpleCPUOffloadScheduler._estimate_lazy_target_blocks(
+        kv_cache_config, max_batched, cp_world_size=cp_world_size
+    )
+
+    if cp_world_size == 1:
+        assert target_cp == target_base
+    else:
+        assert target_cp < target_base, (
+            f"cp_world_size={cp_world_size}: target_cp={target_cp} should be "
+            f"less than target_base={target_base}"
+        )
@@ -53,9 +53,39 @@ PEAGLE_CONFIG = SpeculatorTestConfig(
    parallel_drafting=True,
 )

+QWEN3_EAGLE3_CONFIG = SpeculatorTestConfig(
+    model_path=(
+        "inference-optimization/"
+        "Qwen3-8B-from-Qwen3-8B_regen-speculators.eagle3-qwen3arch-ckpt1"
+    ),
+    method="eagle3",
+    display_name="Qwen3 Eagle3",
+    expected_gsm8k_accuracy=0.88,
+    accuracy_rtol=0.05,
+    expected_acceptance_len=2.67,
+    acceptance_len_rtol=0.10,
+    expected_per_pos_acceptance_rates=(0.76, 0.55, 0.36),
+    per_pos_rtol=0.10,
+)
+
+QWEN3_PEAGLE_CONFIG = SpeculatorTestConfig(
+    model_path="inference-optimization/Qwen3-8B-speculators.peagle-qwen3arch-ckpt4",
+    method="eagle3",
+    display_name="Qwen3 PEagle",
+    expected_gsm8k_accuracy=0.88,
+    accuracy_rtol=0.05,
+    expected_acceptance_len=3.42,
+    acceptance_len_rtol=0.15,
+    expected_per_pos_acceptance_rates=(0.78, 0.59, 0.43, 0.29, 0.18, 0.10, 0.05),
+    per_pos_rtol=0.10,
+    parallel_drafting=True,
+)
+
 SPECULATOR_CONFIGS = [
    pytest.param(DFLASH_CONFIG, id="dflash"),
    pytest.param(PEAGLE_CONFIG, id="peagle"),
+    pytest.param(QWEN3_EAGLE3_CONFIG, id="qwen3arch_eagle3"),
+    pytest.param(QWEN3_PEAGLE_CONFIG, id="qwen3arch_peagle"),
 ]


@@ -176,6 +206,7 @@ def test_speculators_correctness(monkeypatch, config):

    results = evaluate_gsm8k_offline(spec_llm)
    accuracy = results["accuracy"]
+    print(f"GSM8K Accuracy: {accuracy:.4f}")
    accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol)
    assert accuracy >= accuracy_threshold, (
        f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"
@@ -35,7 +35,6 @@ def mock_model_runner_with_input_batch():
        max_model_len=1024,
        max_num_batched_tokens=1024,
        device="cpu",
-        pin_memory=False,
        vocab_size=32000,
        block_sizes=[16],
        kernel_block_sizes=[16],
@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVQuantMode
+from vllm.v1.worker.gpu.attn_utils import _reshape_kv_cache
+from vllm.v1.worker.utils import AttentionGroup
+
+
+class FakeFlashAttentionBackend:
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, 2, block_size, num_kv_heads, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 2, 3, 4)
+
+
+class FakeHNDFlashAttentionBackend(FakeFlashAttentionBackend):
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 3, 2, 4)
+
+
+def test_reshape_padded_flash_attention_kv_cache_strides_by_page():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=2,
+        dtype=torch.float32,
+        page_size_padded=384,
+    )
+    assert spec.real_page_size_bytes == 256
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakeFlashAttentionBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "auto",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 2, 16, 1, 2)
+    assert kv_cache.stride(0) == spec.page_size_bytes // 4
+    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
+    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
+    assert (
+        kv_cache[1, 1].storage_offset()
+        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
+    )
+
+
+def test_reshape_padded_hnd_flash_attention_kv_cache_strides_by_page():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=3,
+        head_size=2,
+        dtype=torch.float32,
+        page_size_padded=1024,
+    )
+    assert spec.real_page_size_bytes == 768
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakeHNDFlashAttentionBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "auto",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 2, 16, 3, 2)
+    assert kv_cache.stride(0) == spec.page_size_bytes // 4
+    assert kv_cache.stride(1) == spec.real_page_size_bytes // 2 // 4
+    assert kv_cache.stride(2) == 2
+    assert kv_cache.stride(3) == spec.block_size * spec.head_size
+    assert kv_cache[1, 0].storage_offset() == spec.page_size_bytes // 4
+    assert (
+        kv_cache[1, 1].storage_offset()
+        == (spec.page_size_bytes + spec.real_page_size_bytes // 2) // 4
+    )
+    assert (
+        kv_cache[1, 1, 3, 2].storage_offset()
+        == (
+            spec.page_size_bytes
+            + spec.real_page_size_bytes // 2
+            + 3 * spec.head_size * 4
+            + 2 * spec.block_size * spec.head_size * 4
+        )
+        // 4
+    )
+
+
+class FakeDiffKVBackend:
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, block_size, num_kv_heads, head_size * 2)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 2, 3)
+
+
+def test_reshape_padded_diff_kv_cache_does_not_infer_kv_dim():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=2,
+        dtype=torch.float32,
+        page_size_padded=384,
+    )
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakeDiffKVBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "auto",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 16, 1, 4)
+    assert kv_cache.stride(0) == spec.page_size_bytes // 4
+    assert kv_cache.stride(1) == 4
+
+
+class FakePerTokenScaleBackend:
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        return (num_blocks, 2, block_size, num_kv_heads, head_size + 4)
+
+    @staticmethod
+    def get_kv_cache_stride_order(
+        include_num_layers_dimension: bool = False,
+    ) -> tuple[int, ...]:
+        assert not include_num_layers_dimension
+        return (0, 1, 2, 3, 4)
+
+
+def test_reshape_padded_quantized_kv_cache_preserves_scale_stride():
+    num_blocks = 3
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=4,
+        dtype=torch.int8,
+        kv_quant_mode=KVQuantMode.INT8_PER_TOKEN_HEAD,
+        page_size_padded=384,
+    )
+    assert spec.real_page_size_bytes == 128
+    assert spec.page_size_bytes == 384
+
+    raw_tensors = {
+        "layer": torch.zeros(spec.page_size_bytes * num_blocks, dtype=torch.int8)
+    }
+    attn_groups = [
+        AttentionGroup(
+            backend=FakePerTokenScaleBackend,
+            layer_names=["layer"],
+            kv_cache_spec=spec,
+            kv_cache_group_id=0,
+        )
+    ]
+
+    kv_cache = _reshape_kv_cache(
+        attn_groups,
+        raw_tensors,
+        "int8_per_token_head",
+        [spec.block_size],
+        {},
+    )["layer"]
+
+    assert kv_cache.shape == (num_blocks, 2, 16, 1, 8)
+    assert kv_cache.stride(0) == spec.page_size_bytes
+    assert kv_cache.stride(1) == 16 * 1 * 8
+    assert kv_cache[1, 1].storage_offset() == spec.page_size_bytes + 16 * 1 * 8
@@ -10,7 +10,6 @@ import torch

 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
-from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.utils.torch_utils import make_tensor_with_pad
 from vllm.v1.pool.metadata import PoolingMetadata
 from vllm.v1.sample.logits_processor import LogitsProcessors
@@ -236,7 +235,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
        max_model_len=1024,
        max_num_batched_tokens=1024,
        device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
        vocab_size=1024,
        block_sizes=[1],
        kernel_block_sizes=[1],
@@ -331,7 +329,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
        max_model_len=1024,
        max_num_batched_tokens=1024,
        device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
        vocab_size=1024,
        block_sizes=[1],
        kernel_block_sizes=[1],
@@ -341,7 +338,6 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
        max_model_len=1024,
        max_num_batched_tokens=1024,
        device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
        vocab_size=1024,
        block_sizes=[1],
        kernel_block_sizes=[1],
@@ -410,7 +406,6 @@ def test_pooling_prompt_lens_not_aliased(device: str):
        max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
        max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS),
        device=torch.device(device),
-        pin_memory=is_pin_memory_available(),
        vocab_size=VOCAB_SIZE,
        block_sizes=[16],
        kernel_block_sizes=[16],
@@ -459,7 +454,6 @@ def test_pooling_metadata_token_id_buffers(
        max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
        max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
        device=torch.device("cpu"),
-        pin_memory=False,
        vocab_size=VOCAB_SIZE,
        block_sizes=[16],
        kernel_block_sizes=[16],
@@ -85,7 +85,6 @@ def initialize_kv_cache(runner: GPUModelRunner):
        max_model_len=runner.max_model_len,
        max_num_batched_tokens=runner.max_num_tokens,
        device=runner.device,
-        pin_memory=runner.pin_memory,
        vocab_size=runner.model_config.get_vocab_size(),
        block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
        kernel_block_sizes=[
@@ -1405,7 +1404,6 @@ def test_input_batch_with_kernel_block_sizes():
    max_model_len = 512
    max_num_batched_tokens = 512
    device = torch.device(DEVICE_TYPE)
-    pin_memory = False
    vocab_size = 50272

    # Test with different kernel block sizes
@@ -1417,7 +1415,6 @@ def test_input_batch_with_kernel_block_sizes():
        max_model_len=max_model_len,
        max_num_batched_tokens=max_num_batched_tokens,
        device=device,
-        pin_memory=pin_memory,
        vocab_size=vocab_size,
        block_sizes=block_sizes,
        kernel_block_sizes=kernel_block_sizes,
@@ -1478,7 +1475,6 @@ def test_hybrid_cache_integration(default_vllm_config, dist_init):
        max_model_len=runner.max_model_len,
        max_num_batched_tokens=runner.max_num_tokens,
        device=runner.device,
-        pin_memory=runner.pin_memory,
        vocab_size=runner.model_config.get_vocab_size(),
        block_sizes=[kv_cache_config.kv_cache_groups[0].kv_cache_spec.block_size],
        kernel_block_sizes=[16],
@@ -991,7 +991,7 @@ class VllmBackend:
            },
            payload_fn=lambda: json.dumps(
                {
-                    "model": self.vllm_config.model_config.model,
+                    "model": getattr(self.vllm_config.model_config, "model", "unknown"),
                    "prefix": self.prefix,
                    "mode": str(cc.mode),
                    "backend": cc.backend,
@@ -302,6 +302,14 @@ class ParallelConfig:
    Each entry must use `numactl --physcpubind` CPU-list syntax, for example
    `"0-3"` or `"0,2,4-7"`.
    """
+    assigned_physical_gpu_ids: list[int] | None = None
+    """Mapping from vLLM-local logical GPU IDs to physical GPU IDs.
+
+    For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2,
+    and logical GPU 1 maps to physical GPU 3. Physical IDs are used only
+    at platform/topology boundaries such as NVML, NIC affinity, P2P
+    checks, and final CUDA device selection when needed. When None,
+    logical IDs map to visible device IDs in order."""

    distributed_timeout_seconds: int | None = None
    """Timeout in seconds for distributed operations (e.g., init_process_group).
@@ -772,6 +780,7 @@ class ParallelConfig:
            "numa_bind",
            "numa_bind_nodes",
            "numa_bind_cpus",
+            "assigned_physical_gpu_ids",
        }

        from vllm.config.utils import get_hash_factors, hash_factors
@@ -18,8 +18,8 @@ import torch

 from vllm.device_allocator import AllocationData, HandleType
 from vllm.logger import init_logger
-from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.utils.system_utils import find_loaded_library
+from vllm.utils.torch_utils import PIN_MEMORY

 logger = init_logger(__name__)

@@ -196,7 +196,7 @@ class CuMemAllocator:
                    size_in_bytes,
                    dtype=torch.uint8,
                    device="cpu",
-                    pin_memory=is_pin_memory_available(),
+                    pin_memory=PIN_MEMORY,
                )
                cpu_ptr = cpu_backup_tensor.data_ptr()
                libcudart.cudaMemcpy(cpu_ptr, ptr, size_in_bytes)
@@ -11,7 +11,7 @@ import torch

 from vllm.device_allocator import AllocationData, HandleType
 from vllm.logger import init_logger
-from vllm.utils.platform_utils import is_pin_memory_available
+from vllm.utils.torch_utils import PIN_MEMORY

 logger = init_logger(__name__)

@@ -188,7 +188,7 @@ class XpuMemAllocator:
                size_in_bytes,
                dtype=torch.uint8,
                device="cpu",
-                pin_memory=is_pin_memory_available(),
+                pin_memory=PIN_MEMORY,
            )
            cpu_ptr = cpu_backup_tensor.data_ptr()
            _xpu_memcpy_sync(
@@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
        self.num_experts = num_experts

        self.cleanup()
-        gpus_per_node = torch.accelerator.device_count()
+        from vllm.platforms.interface import get_assigned_physical_gpu_ids
+
+        assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
+        gpus_per_node = (
+            len(assigned_physical_gpu_ids)
+            if assigned_physical_gpu_ids is not None
+            else torch.accelerator.device_count()
+        )
        logger.debug(
            "Making One-sided NVLink mapping: rank=%d, world size=%d",
            self.rank,
@@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:

    is_distributed = dist.is_initialized()

-    num_dev = current_platform.device_count()
-    cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
-    if cuda_visible_devices is None:
-        cuda_visible_devices = ",".join(str(i) for i in range(num_dev))
+    from vllm.platforms.interface import get_assigned_physical_gpu_ids
+
+    assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
+    if assigned_physical_gpu_ids is not None:
+        # Key by the ordered list: the cache stores directed local-index
+        # pairs, so permutations of the same set are distinct mappings.
+        cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids)
+        num_dev = len(assigned_physical_gpu_ids)
+    else:
+        num_dev = current_platform.device_count()
+        cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
+        cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev))

    path = os.path.join(
-        envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json"
+        envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json"
    )
    os.makedirs(os.path.dirname(path), exist_ok=True)
    from vllm.distributed.parallel_state import get_world_group
@@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
        #  enter this block to calculate the cache
        logger.info("generating GPU P2P access cache in %s", path)
        cache: dict[str, bool] = {}
-        ids = list(range(num_dev))
+        # The probe subprocesses inherit this process's device-control env
+        # var, so they must be given visible ordinals, not physical IDs.
+        if assigned_physical_gpu_ids is not None:
+            ids = [
+                current_platform.logical_device_id_to_visible_device_id(local)
+                for local in range(num_dev)
+            ]
+        else:
+            ids = list(range(num_dev))
        # batch of all pairs of GPUs
        batch_src, batch_tgt = zip(*list(product(ids, ids)))
        # NOTE: we use `subprocess` rather than `multiprocessing` here
@@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
                ) from e
            with open(output_file.name, "rb") as f:
                result = pickle.load(f)
+        # Cache entries must be keyed by local indices (0..N-1) because
+        # gpu_p2p_access_check() is called with local ranks.
+        id_to_local = {device_id: local for local, device_id in enumerate(ids)}
        for _i, _j, r in zip(batch_src, batch_tgt, result):
-            cache[f"{_i}->{_j}"] = r
+            cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r
        with open(path, "w") as f:
            json.dump(cache, f, indent=4)
    if is_distributed:
@@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool:
            continue
        if envs.VLLM_SKIP_P2P_CHECK:
            logger.debug("Skipping P2P check and trusting the driver's P2P report.")
-            return torch.cuda.can_device_access_peer(rank, i)
+            # can_device_access_peer takes visible device ordinals, while
+            # rank and i are logical local IDs.
+            return torch.cuda.can_device_access_peer(
+                current_platform.logical_device_id_to_visible_device_id(rank),
+                current_platform.logical_device_id_to_visible_device_id(i),
+            )
        if not gpu_p2p_access_check(rank, i):
            return False
    return True
@@ -126,13 +131,10 @@ class CustomAllreduce:
                    CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
                    max_size,
                )
-        cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
-        if cuda_visible_devices:
-            device_ids = list(map(int, cuda_visible_devices.split(",")))
-        else:
-            device_ids = list(range(current_platform.device_count()))
-
-        physical_device_id = device_ids[device.index]
+        # device.index is a visible ordinal, not a logical local ID.
+        physical_device_id = current_platform.visible_device_id_to_physical_device_id(
+            device.index
+        )
        tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
        gather_list = [
            torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
@@ -129,12 +129,10 @@ class QuickAllReduce:
        assert isinstance(device, torch.device)
        self.device = device

-        cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
-        if cuda_visible_devices:
-            device_ids = list(map(int, cuda_visible_devices.split(",")))
-        else:
-            device_ids = list(range(current_platform.device_count()))
-        physical_device_id = device_ids[device.index]
+        # device.index is a visible ordinal, not a logical local ID.
+        physical_device_id = current_platform.visible_device_id_to_physical_device_id(
+            device.index
+        )
        tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
        gather_list = [
            torch.tensor([0], dtype=torch.int, device="cpu")
@@ -840,7 +840,13 @@ class MessageQueue:
            The MessageQueue instance for the calling process,
            and a list of handles (only non-empty for the reader process).
        """
-        local_size = current_platform.device_count()
+        from vllm.platforms.interface import get_assigned_physical_gpu_ids
+
+        assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
+        if assigned_physical_gpu_ids is not None:
+            local_size = len(assigned_physical_gpu_ids)
+        else:
+            local_size = current_platform.device_count()
        rank = dist.get_rank()
        same_node = rank // local_size == reader_rank // local_size
        buffer_io = MessageQueue(
@@ -482,10 +482,11 @@ def _init_lmcache_engine(
    )

    # Change current device.
-    num_gpus = torch.accelerator.device_count()
-    local_rank = parallel_config.rank % num_gpus
-    torch.accelerator.set_device_index(local_rank)
-    device = torch.device(f"cuda:{local_rank}")
+    from vllm.distributed.parallel_state import get_world_group
+
+    device_index = get_world_group().device_index
+    torch.accelerator.set_device_index(device_index)
+    device = torch.device(f"cuda:{device_index}")
    metadata = LMCacheEngineMetadata(
        model_config.model,
        parallel_config.world_size,
@@ -2,13 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """External-store cache-hit coordinator for MooncakeStoreConnector."""

+from collections.abc import Sequence
 from typing import cast

+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
+    chunk_hashes_for_block_size,
+)
 from vllm.v1.core.block_pool import BlockPool
 from vllm.v1.core.kv_cache_utils import (
    BlockHash,
-    BlockHashList,
-    BlockHashListWithBlockSize,
    KVCacheBlock,
 )
 from vllm.v1.core.single_type_kv_cache_manager import (
@@ -120,7 +122,7 @@ class MooncakeStoreCoordinator:

    def find_longest_cache_hit(
        self,
-        block_hashes: list[BlockHash],
+        block_hashes: Sequence[BlockHash],
        max_length: int,
        cached_block_pool: ExternalCachedBlockPool,
        *,
@@ -147,7 +149,7 @@ class MooncakeStoreCoordinator:

    def load_mask(
        self,
-        block_hashes: list[BlockHash],
+        block_hashes: Sequence[BlockHash],
        token_len: int,
    ) -> tuple[list[bool], ...]:
        """Per-group load masks: ``mask[g][i]`` is True iff group ``g``'s
@@ -236,17 +238,15 @@ class MooncakeStoreCoordinator:
        return tuple(masks)

    def block_hashes_for_spec(
-        self, block_hashes: list[BlockHash], spec: KVCacheSpec
-    ) -> BlockHashList:
-        if spec.block_size == self.hash_block_size:
-            return block_hashes
-        return BlockHashListWithBlockSize(
+        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: list[BlockHash],
+        block_hashes: Sequence[BlockHash],
        max_length: int,
        cached_block_pool: ExternalCachedBlockPool,
        *,
@@ -264,7 +264,7 @@ class MooncakeStoreCoordinator:
            spec, group_ids, manager_cls = self.attention_groups[0]
            hashes = self.block_hashes_for_spec(block_hashes, spec)
            hit_blocks = manager_cls.find_longest_cache_hit(
-                block_hashes=hashes,
+                block_hashes=hashes,  # type: ignore[arg-type]
                max_length=max_length,
                kv_cache_group_ids=group_ids,
                block_pool=cast(BlockPool, cached_block_pool),
@@ -304,7 +304,7 @@ class MooncakeStoreCoordinator:
                    _max_length = min(curr_hit_length + spec.block_size, max_length)
                hashes = self.block_hashes_for_spec(block_hashes, spec)
                hit_blocks = manager_cls.find_longest_cache_hit(
-                    block_hashes=hashes,
+                    block_hashes=hashes,  # type: ignore[arg-type]
                    max_length=_max_length,
                    kv_cache_group_ids=group_ids,
                    block_pool=cast(BlockPool, cached_block_pool),
@@ -5,8 +5,9 @@
 # (vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/).
 """Data classes for MooncakeStoreConnector."""

-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
+from typing import cast

 import torch

@@ -23,6 +24,77 @@ from vllm.v1.core.kv_cache_utils import (
 logger = init_logger(__name__)


+class BlobBlockHashes(Sequence[BlockHash]):
+    """Lazy view over a flat buffer of fixed-size block hashes to avoid the overhead
+    of materializing all hashes upfront.
+    """
+
+    def __init__(self, blob: memoryview, hash_len: int):
+        self._blob = blob
+        self._hash_len = hash_len
+        self._n = len(blob) // hash_len if hash_len else 0
+
+    def __len__(self) -> int:
+        return self._n
+
+    def __getitem__(self, idx):
+        if isinstance(idx, slice):
+            return [self[i] for i in range(*idx.indices(self._n))]
+        if idx < 0:
+            idx += self._n
+        if not 0 <= idx < self._n:
+            raise IndexError(idx)
+        off = idx * self._hash_len
+        return BlockHash(self._blob[off : off + self._hash_len])
+
+
+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``).
+    """
+
+    def __init__(
+        self,
+        block_hashes: Sequence[BlockHash],
+        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.
+        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
+
+    def _get_value_at(self, idx: int) -> BlockHash:
+        return self.block_hashes[idx * self.scale_factor + self.scale_factor - 1]
+
+
+def chunk_hashes_for_block_size(
+    block_hashes: Sequence[BlockHash],
+    hash_block_size: int,
+    block_size: int,
+) -> Sequence[BlockHash]:
+    """Map ``hash_block_size``-granular block hashes to one compact hash per
+    ``block_size`` chunk (the chunk's last sub-hash). Returns ``block_hashes``
+    unchanged when the two sizes are equal.
+    """
+    if block_size == hash_block_size:
+        return block_hashes
+    # Structurally a Sequence[BlockHash] (indexable + sized); the base class
+    # just isn't declared as one.
+    return cast(
+        "Sequence[BlockHash]",
+        _CompactChunkHashList(block_hashes, hash_block_size, block_size),
+    )
+
+
@dataclass
 class KeyMetadata:
    """Metadata for constructing pool keys."""
@@ -138,18 +210,15 @@ class ChunkedTokenDatabase:
        Args:
            token_len: Total number of tokens.
            block_hashes: Block hashes computed at ``hash_block_size`` granularity.
-                When ``block_size > hash_block_size`` consecutive hashes are merged
-                up to the group's ``block_size`` via ``BlockHashListWithBlockSize``.
+                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``.
            mask_num: Number of tokens to skip from the beginning.
        """
        if not block_hashes:
            return
-        if self.block_size == self.hash_block_size:
-            chunk_hashes: Iterable[BlockHash] = block_hashes
-        else:
-            chunk_hashes = BlockHashListWithBlockSize(
-                block_hashes, self.hash_block_size, self.block_size
-            )
+        chunk_hashes: Iterable[BlockHash] = chunk_hashes_for_block_size(
+            block_hashes, self.hash_block_size, self.block_size
+        )
        for chunk_id, h in enumerate(chunk_hashes):
            start_idx = chunk_id * self.block_size
            if start_idx >= token_len:
@@ -11,7 +11,10 @@ Wire format (REQ/REP over IPC):

      msg_type == LOOKUP_MSG:
          frame 1: token_len (u32 big-endian, 4 bytes)
-          frame 2..n: msgpack-encoded list[str] of block-hash hex digests
+          frame 2: hash_len (u16 big-endian, 2 bytes) — byte length of each
+                   fixed-size block hash (0 when there are no hashes)
+          frame 3: raw block hashes concatenated back-to-back (each hash_len
+                   bytes); the server splits on hash_len
        Response: [hit_count: u32 big-endian, 4 bytes]

      msg_type == RESET_MSG:
@@ -18,7 +18,7 @@ import socket
 import threading
 import time
 from collections import defaultdict
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import Any, Literal, TypeVar
@@ -45,6 +45,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.coordinator imp
    MooncakeStoreCoordinator,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (  # noqa: E501
+    BlobBlockHashes,
    ChunkedTokenDatabase,
    KeyMetadata,
    MooncakeStoreConnectorMetadata,
@@ -65,7 +66,6 @@ from vllm.v1.core.kv_cache_utils import (
    resolve_kv_cache_block_sizes,
 )
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec
-from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder

 from .metrics import MooncakeStoreConnectorStats

@@ -1372,7 +1372,7 @@ class MooncakeStoreWorker:

        return finished_sending

-    def lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
+    def lookup(self, token_len: int, block_hashes: Sequence[BlockHash]) -> int:
        """Check how many prefix tokens exist in the store.

        Checks across all TP ranks and PP ranks.
@@ -1392,6 +1392,11 @@ class MooncakeStoreWorker:
            group_hashes = self.coord.block_hashes_for_spec(
                block_hashes, self._kv_cache_groups[g_idx].kv_cache_spec
            )
+            metadata_templates = [
+                dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
+                for tp in range(tp_count)
+                for pp in range(self.pp_size)
+            ]
            for chunk_id, h in enumerate(group_hashes):
                start_idx = chunk_id * spec_block_size
                if start_idx >= token_len:
@@ -1400,11 +1405,11 @@ class MooncakeStoreWorker:
                    chunk_id >= len(lookup_mask) or not lookup_mask[chunk_id]
                ):
                    continue
-                for tp in range(tp_count):
-                    for pp in range(self.pp_size):
-                        md = dataclasses.replace(db.metadata, tp_rank=tp, pp_rank=pp)
-                        candidate_keys.append(PoolKey(md, h.hex()).to_string())
-                        candidate_meta.append((g_idx, bytes(h)))
+                h_hex = h.hex()
+                h_bytes = bytes(h)
+                for md in metadata_templates:
+                    candidate_keys.append(PoolKey(md, h_hex).to_string())
+                    candidate_meta.append((g_idx, h_bytes))

        if not candidate_keys:
            return 0
@@ -1483,7 +1488,6 @@ class LookupKeyServer:
        store_worker: MooncakeStoreWorker,
        vllm_config: VllmConfig,
    ):
-        self.decoder = MsgpackDecoder()
        self.ctx = zmq.Context()  # type: ignore[attr-defined]
        socket_path = get_zmq_rpc_path_lookup(vllm_config)
        self._ipc_path = socket_path.removeprefix("ipc://")
@@ -1506,9 +1510,9 @@ class LookupKeyServer:

                if msg_type == LOOKUP_MSG:
                    token_len = int.from_bytes(all_frames[1], byteorder="big")
-                    hash_frames = all_frames[2:]
-                    hashes_str = self.decoder.decode(hash_frames)
-                    block_hashes = [BlockHash(bytes.fromhex(s)) for s in hashes_str]
+                    hash_len = int.from_bytes(all_frames[2], byteorder="big")
+                    blob = all_frames[3].buffer
+                    block_hashes = BlobBlockHashes(blob, hash_len)
                    result = self.store_worker.lookup(token_len, block_hashes)
                    self.socket.send(result.to_bytes(4, "big"))

@@ -1557,7 +1561,6 @@ class LookupKeyClient:
    """

    def __init__(self, vllm_config: VllmConfig):
-        self.encoder = MsgpackEncoder()
        self.ctx = zmq.Context()  # type: ignore[attr-defined]
        socket_path = get_zmq_rpc_path_lookup(vllm_config)
        self.socket = make_zmq_socket(
@@ -1574,14 +1577,16 @@ class LookupKeyClient:
        self.futures: dict[str, Future[int]] = {}

    def _lookup(self, token_len: int, block_hashes: list[BlockHash]) -> int:
-        hash_strs = [h.hex() for h in block_hashes]
-        hash_frames = self.encoder.encode(hash_strs)
-        token_len_bytes = token_len.to_bytes(4, byteorder="big")
-        all_frames = [LOOKUP_MSG, token_len_bytes] + list(hash_frames)
+        hash_len = len(block_hashes[0]) if block_hashes else 0
+        all_frames = (
+            LOOKUP_MSG,
+            token_len.to_bytes(4, byteorder="big"),
+            hash_len.to_bytes(2, byteorder="big"),
+            b"".join(block_hashes),
+        )
        self.socket.send_multipart(all_frames, copy=False)
        resp = self.socket.recv()
-        result = int.from_bytes(resp, "big")
-        return result
+        return int.from_bytes(resp, "big")

    def lookup(
        self,
@@ -47,6 +47,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import (
    MoRIIOWrapper,
    MoRIIOWriter,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import (
+    LayerTransferGeometry,
+    build_layer_to_spec,
+    compute_block_transfer_offsets,
+    get_layer_transfer_geometry,
+    is_mla_cache_layer,
+    iter_layer_registration_regions,
+)
 from vllm.distributed.parallel_state import (
    get_tensor_model_parallel_world_size,
    get_tp_group,
@@ -71,6 +79,7 @@ if TYPE_CHECKING:

 logger = init_logger(__name__)

+
 try:
    from mori.io import (
        BackendType,
@@ -117,7 +126,9 @@ class MoRIIOConnector(KVConnectorBase_V1):
            self.connector_worker: MoRIIOConnectorWorker | None = None
        elif role == KVConnectorRole.WORKER:
            self.connector_scheduler = None
-            self.connector_worker = MoRIIOConnectorWorker(vllm_config, self.engine_id)
+            self.connector_worker = MoRIIOConnectorWorker(
+                vllm_config, self.engine_id, kv_cache_config
+            )
        logger.info(
            "Initialized MoRIIO Connector,engine_id:%s,role: %s",
            self.engine_id,
@@ -683,7 +694,12 @@ class MoRIIOConnectorScheduler:
 class MoRIIOConnectorWorker:
    """Implementation of Worker side methods"""

-    def __init__(self, vllm_config: VllmConfig, engine_id: str):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        engine_id: str,
+        kv_cache_config: "KVCacheConfig",
+    ):
        if not is_moriio_available():
            raise RuntimeError(
                "MoRIIO is not available. Please ensure the 'mori' package "
@@ -707,6 +723,7 @@ class MoRIIOConnectorWorker:
        )
        self.kv_transfer_config = vllm_config.kv_transfer_config
        self.is_producer = self.kv_transfer_config.is_kv_producer
+        self.layer_to_spec = build_layer_to_spec(kv_cache_config)

        if self.is_producer:
            set_role(ROLE.PRODUCER)
@@ -809,6 +826,8 @@ class MoRIIOConnectorWorker:
        self.kv_cache_shape = None
        self.block_shape = None
        self.kv_element_size = 0
+        self.kv_cache_shapes: dict[str, torch.Size] = {}
+        self.block_lens: dict[str, int] = {}

        # Map of engine_id -> {agent_name0, agent_name1..}.
        self._remote_agents: dict[EngineId, set[str]] = {}
@@ -1218,51 +1237,86 @@ class MoRIIOConnectorWorker:
        all_done_future = self._handshake_initiation_executor.submit(wait_all_dp)
        all_done_future.add_done_callback(request_ready)

+    def _is_mla_cache_layer(self, layer_name: str) -> bool:
+        return is_mla_cache_layer(self.layer_to_spec, layer_name)
+
+    def _get_layer_transfer_geometry(
+        self, layer_name: str, remote_num_blocks: int | None = None
+    ) -> LayerTransferGeometry:
+        return get_layer_transfer_geometry(
+            layer_name,
+            self.kv_caches[layer_name],
+            self.layer_to_spec,
+            remote_num_blocks,
+        )
+
+    def _iter_layer_registration_regions(
+        self, layer_name: str
+    ) -> list[tuple[torch.Tensor, int]]:
+        return iter_layer_registration_regions(
+            layer_name,
+            self.kv_caches[layer_name],
+            self.layer_to_spec,
+        )
+
    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
        """Register the KV Cache data in moriio."""

-        _, first_kv_cache = next(iter(kv_caches.items()))
+        self.kv_caches = kv_caches  # layer name to kv cache
+        self.kv_cache_shapes = {
+            layer_name: kv_cache.shape for layer_name, kv_cache in kv_caches.items()
+        }
+
+        first_layer_name, first_kv_cache = next(
+            (
+                (layer_name, kv_cache)
+                for layer_name, kv_cache in kv_caches.items()
+                if (
+                    not self._is_mla_cache_layer(layer_name)
+                    and len(kv_cache.shape) == 5
+                    and (kv_cache.shape[0] == 2 or kv_cache.shape[1] == 2)
+                )
+            ),
+            next(iter(kv_caches.items())),
+        )
        kv_elem_size = first_kv_cache.element_size()

-        use_mla = len(first_kv_cache.shape) == 3
-        assert use_mla == self.use_mla
+        use_mla = self._is_mla_cache_layer(first_layer_name)
+        first_geometry = self._get_layer_transfer_geometry(first_layer_name)

        if use_mla:
            # MLA case.
-            self.num_blocks = first_kv_cache.shape[0]
            block_rank = 2  # [block_size, latent_dim]
            block_shape = first_kv_cache.shape[-block_rank:]
-            block_size, kv_latent_dim = block_shape
-            self.slot_size_bytes = kv_elem_size * kv_latent_dim
        else:
-            # [2 (k and v), num_blocks, ...]
-            self.num_blocks = first_kv_cache.shape[1]
+            # [2, num_blocks, ...] or [num_blocks, 2, ...]
            block_rank = 3  # [block_size, kv_heads, head_dim]
            block_shape = first_kv_cache.shape[-block_rank:]
-            block_size, n_kv_heads, head_dim = block_shape[-3:]
-            # head size in bytes.
-            self.slot_size_bytes = (
-                kv_elem_size * n_kv_heads * head_dim
-            )  # 1 token 1 layer size , slot size
-        assert block_size == self.block_size
+        self.num_blocks = first_geometry.num_blocks
+        self.slot_size_bytes = first_geometry.slot_size_bytes
+        assert first_geometry.block_size == self.block_size
        # TODO(tms): self.block_len needs to be per-layer for sliding window,
        # hybrid attn, etc
        # block size in bytes
-        self.block_len = kv_elem_size * math.prod(block_shape)
+        self.block_len = first_geometry.block_len
        self.kv_cache_shape = first_kv_cache.shape
        self.block_shape = block_shape
        self.kv_element_size = kv_elem_size

        self.dst_num_blocks[self.engine_id] = self.num_blocks
-        self.kv_caches = kv_caches  # layer name to kv cache
        kv_caches_base_addr = []
        caches_data = []

-        for cache_or_caches in kv_caches.values():
-            cache_list = [cache_or_caches] if use_mla else cache_or_caches
-            for cache in cache_list:
+        for layer_name in kv_caches:
+            geometry = self._get_layer_transfer_geometry(layer_name)
+            if geometry.block_size != self.block_size:
+                raise ValueError(
+                    "MoRIIO KV cache block size mismatch for layer "
+                    f"{layer_name}: {geometry.block_size} != {self.block_size}"
+                )
+            self.block_lens[layer_name] = geometry.block_len
+            for cache, region_len in self._iter_layer_registration_regions(layer_name):
                base_addr = cache.data_ptr()
-                region_len = self.num_blocks * self.block_len
                caches_data.append((base_addr, region_len, cache.device.index, ""))
                kv_caches_base_addr.append(base_addr)

@@ -1275,7 +1329,9 @@ class MoRIIOConnectorWorker:
                moriio_mem_metadata
            )

-            self.local_kv_cache_size.append(cache.nelement() * cache.element_size())
+            self.local_kv_cache_size.append(
+                kv_cache.nelement() * kv_cache.element_size()
+            )

        self.kv_caches_base_addr[self.engine_id] = kv_caches_base_addr
        self.num_regions = len(caches_data)
@@ -1666,47 +1722,17 @@ class MoRIIOConnectorWorker:
        Returns:
            Tuple of (local_offsets, remote_offsets, transfer_sizes)
        """
-        assert self.kv_cache_shape is not None, "KV caches shape not initialized"
-        is_mla = len(self.kv_cache_shape) == 3
-        stride = self.kv_caches[layer_name].stride()
-        sz = self.kv_caches[layer_name].element_size()
-        if is_mla:
-            blknum, blksize, hs = self.kv_cache_shape
-            hn = 1
-            block_stride = stride[0]
-        else:
-            _, blknum, blksize, hn, hs = self.kv_cache_shape
-            local_ktov_stride = stride[0]
-            block_stride = stride[1]
-            remote_ktov_stride = block_stride * remote_moriio_meta.num_blocks
-
-        transfer_size_byte = blksize * hn * hs * sz
-        per_block = 1 if is_mla else 2
-        total = len(local_block_ids) * per_block
-        offset_local = [0] * total
-        offset_remote = [0] * total
-        sizes = [transfer_size_byte] * total
-
-        w = 0
-        for i, lb in enumerate(local_block_ids):
-            rb = remote_block_ids[i]
-            # K
-            offset_local[w] = sz * (lb * block_stride)
-            offset_remote[w] = sz * (rb * block_stride)
-            w += 1
-            if not is_mla:
-                # V
-                # Handle num_block variations originating from PD (different kv strides)
-                # TODO: address block_sz differences in heterogeneous TP scenarios
-                # In MLA, we don't need to consider these two cases.
-                offset_local[w] = sz * (1 * local_ktov_stride + lb * block_stride)
-                offset_remote[w] = sz * (1 * remote_ktov_stride + rb * block_stride)
-                w += 1
-
-        merged_l, merged_r, merged_s = self.merge_contiguous_blocks(
-            offset_local, offset_remote, sizes, assume_sorted=False
+        return compute_block_transfer_offsets(
+            layer_name=layer_name,
+            kv_cache=self.kv_caches[layer_name],
+            layer_to_spec=self.layer_to_spec,
+            local_block_ids=local_block_ids,
+            remote_block_ids=remote_block_ids,
+            remote_num_blocks=remote_moriio_meta.num_blocks,
+            merge_fn=lambda local, remote, sizes: self.merge_contiguous_blocks(
+                local, remote, sizes, assume_sorted=False
+            ),
        )
-        return merged_l, merged_r, merged_s

    def _read_blocks(
        self,
@@ -1724,15 +1750,13 @@ class MoRIIOConnectorWorker:
        dp0_engine_id = self.get_engine_name_with_dp(dst_engine_id, 0)
        sessions, remote_moriio_meta = self._get_built_session(dp0_engine_id)

-        first_layer = list(self.layer_name_to_local_kv_cache_metadata.keys())[0]
-        offs = self._compute_block_transfer_offsets(
-            first_layer, local_block_ids, remote_block_ids, remote_moriio_meta
-        )
-
        for layer_name in self.layer_name_to_local_kv_cache_metadata:
            sess_idx = list(self.layer_name_to_local_kv_cache_metadata.keys()).index(
                layer_name
            )
+            offs = self._compute_block_transfer_offsets(
+                layer_name, local_block_ids, remote_block_ids, remote_moriio_meta
+            )
            # TODO : apply multi-session batch-read when moriio support it
            transfer_status = self.moriio_wrapper.read_remote_data(
                offs[2], offs[0], offs[1], sessions[sess_idx]
@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable, Mapping
+from typing import NamedTuple
+
+import torch
+
+from vllm.v1.kv_cache_interface import (
+    KVCacheConfig,
+    KVCacheSpec,
+    MLAAttentionSpec,
+    SlidingWindowMLASpec,
+    UniformTypeKVCacheSpecs,
+)
+
+
+class LayerTransferGeometry(NamedTuple):
+    num_blocks: int
+    block_size: int
+    block_len: int
+    slot_size_bytes: int
+    block_stride: int
+    local_kv_stride: int | None
+    remote_kv_stride: int | None
+    transfers_per_block: int
+    regions_per_block: int
+    split_kv_regions: bool
+
+
+def build_layer_to_spec(kv_cache_config: KVCacheConfig) -> dict[str, KVCacheSpec]:
+    layer_to_spec: dict[str, KVCacheSpec] = {}
+    for group in kv_cache_config.kv_cache_groups:
+        group_spec = group.kv_cache_spec
+        if isinstance(group_spec, UniformTypeKVCacheSpecs):
+            layer_to_spec.update(
+                {
+                    layer_name: group_spec.kv_cache_specs[layer_name]
+                    for layer_name in group.layer_names
+                }
+            )
+        else:
+            layer_to_spec.update(
+                {layer_name: group_spec for layer_name in group.layer_names}
+            )
+    return layer_to_spec
+
+
+def is_mla_cache_layer(
+    layer_to_spec: Mapping[str, KVCacheSpec], layer_name: str
+) -> bool:
+    try:
+        spec = layer_to_spec[layer_name]
+    except KeyError as e:
+        raise ValueError(f"Missing KV cache spec for layer {layer_name}") from e
+    return isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
+
+
+def get_layer_transfer_geometry(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+    remote_num_blocks: int | None = None,
+) -> LayerTransferGeometry:
+    shape = kv_cache.shape
+    stride = kv_cache.stride()
+    element_size = kv_cache.element_size()
+    is_mla_cache = is_mla_cache_layer(layer_to_spec, layer_name)
+
+    if is_mla_cache and len(shape) == 3:
+        num_blocks, block_size, latent_dim = shape
+        slot_size_bytes = latent_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[0],
+            local_kv_stride=None,
+            remote_kv_stride=None,
+            transfers_per_block=1,
+            regions_per_block=1,
+            split_kv_regions=False,
+        )
+
+    if not is_mla_cache and len(shape) == 5 and shape[0] == 2:
+        _, num_blocks, block_size, num_kv_heads, head_dim = shape
+        slot_size_bytes = num_kv_heads * head_dim * element_size
+        block_len = block_size * slot_size_bytes
+        remote_kv_stride = stride[1] * (remote_num_blocks or num_blocks)
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[1],
+            local_kv_stride=stride[0],
+            remote_kv_stride=remote_kv_stride,
+            transfers_per_block=2,
+            regions_per_block=1,
+            split_kv_regions=True,
+        )
+
+    if not is_mla_cache and len(shape) == 5 and shape[1] == 2:
+        num_blocks, _, block_size, num_kv_heads, head_dim = shape
+        slot_size_bytes = num_kv_heads * head_dim * element_size
+        block_len = block_size * slot_size_bytes
+        return LayerTransferGeometry(
+            num_blocks=num_blocks,
+            block_size=block_size,
+            block_len=block_len,
+            slot_size_bytes=slot_size_bytes,
+            block_stride=stride[0],
+            local_kv_stride=stride[1],
+            remote_kv_stride=stride[1],
+            transfers_per_block=2,
+            regions_per_block=2,
+            split_kv_regions=False,
+        )
+
+    cache_kind = "MLA" if is_mla_cache else "K/V"
+    raise ValueError(
+        f"Unsupported MoRIIO {cache_kind} cache shape for layer "
+        f"{layer_name}: {tuple(shape)}"
+    )
+
+
+def iter_layer_registration_regions(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+) -> list[tuple[torch.Tensor, int]]:
+    geometry = get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec)
+    region_len = geometry.num_blocks * geometry.regions_per_block * geometry.block_len
+    if geometry.split_kv_regions:
+        return [(cache, region_len) for cache in kv_cache]
+    return [(kv_cache, region_len)]
+
+
+def merge_contiguous_offsets(
+    offsets_local: list[int],
+    offsets_remote: list[int],
+    sizes: list[int],
+) -> tuple[list[int], list[int], list[int]]:
+    if not offsets_local:
+        return [], [], []
+    if not (len(offsets_local) == len(offsets_remote) == len(sizes)):
+        raise ValueError("Input list lengths mismatch")
+
+    rows = sorted(zip(offsets_local, offsets_remote, sizes), key=lambda row: row[0])
+    merged: list[list[int]] = []
+    for local, remote, size in rows:
+        if (
+            merged
+            and local == merged[-1][0] + merged[-1][2]
+            and remote == merged[-1][1] + merged[-1][2]
+        ):
+            merged[-1][2] += size
+        else:
+            merged.append([local, remote, size])
+
+    return (
+        [row[0] for row in merged],
+        [row[1] for row in merged],
+        [row[2] for row in merged],
+    )
+
+
+def compute_block_transfer_offsets(
+    layer_name: str,
+    kv_cache: torch.Tensor,
+    layer_to_spec: Mapping[str, KVCacheSpec],
+    local_block_ids: list[int],
+    remote_block_ids: list[int],
+    remote_num_blocks: int,
+    merge_fn: Callable[
+        [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]]
+    ] = merge_contiguous_offsets,
+) -> tuple[list[int], list[int], list[int]]:
+    if len(local_block_ids) != len(remote_block_ids):
+        raise ValueError(
+            "local_block_ids and remote_block_ids must have the same length: "
+            f"{len(local_block_ids)} != {len(remote_block_ids)}"
+        )
+    geometry = get_layer_transfer_geometry(
+        layer_name, kv_cache, layer_to_spec, remote_num_blocks
+    )
+    element_size = kv_cache.element_size()
+    transfer_size_byte = geometry.block_len
+    per_block = geometry.transfers_per_block
+    total = len(local_block_ids) * per_block
+    offset_local = [0] * total
+    offset_remote = [0] * total
+    sizes = [transfer_size_byte] * total
+
+    w = 0
+    for lb, rb in zip(local_block_ids, remote_block_ids):
+        offset_local[w] = element_size * (lb * geometry.block_stride)
+        offset_remote[w] = element_size * (rb * geometry.block_stride)
+        w += 1
+        if per_block == 2:
+            assert geometry.local_kv_stride is not None
+            assert geometry.remote_kv_stride is not None
+            offset_local[w] = element_size * (
+                geometry.local_kv_stride + lb * geometry.block_stride
+            )
+            offset_remote[w] = element_size * (
+                geometry.remote_kv_stride + rb * geometry.block_stride
+            )
+            w += 1
+
+    return merge_fn(offset_local, offset_remote, sizes)
@@ -841,8 +841,106 @@ class NixlBaseConnectorWorker:
        # Forwarding a real layer name rather than a synthetic key
        self.register_kv_caches({first_layer: kv_cache})

+    def _register_packed_kv_cache(
+        self,
+        storage: torch.UntypedStorage,
+    ) -> None:
+        """Register a packed KV cache as a single NIXL region.
+
+        The packed allocation interleaves all layers per block, so each
+        block_stride-byte chunk is one logical block.  We register 1
+        NIXL region and create 1 descriptor per block.
+        """
+        self.transfer_topo = TransferTopology(
+            tp_rank=self.tp_rank,
+            tp_size=self.world_size,
+            block_size=self.block_size,
+            engine_id=self.engine_id,
+            is_mla=self.use_mla,
+            total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
+            attn_backends=self.attn_backends,
+            tensor_shape=None,
+            is_mamba=self._has_mamba,
+        )
+        self.compat_hash = compute_nixl_compatibility_hash(
+            self.vllm_config,
+            self.backend_name,
+            self.transfer_topo.cross_layers_blocks,
+        )
+
+        total_size = storage.nbytes()
+        block_stride = total_size // self.num_blocks
+        base_addr = storage.data_ptr()
+        device_id = storage.device.index
+        assert device_id is not None
+
+        logger.info(
+            "Registering packed KV cache: total_size=%s, block_stride=%s, "
+            "num_blocks=%s, num_regions=1",
+            total_size,
+            block_stride,
+            self.num_blocks,
+        )
+
+        self.device_id = device_id
+        caches_data = [(base_addr, total_size, self.device_id, "")]
+
+        self.block_len_per_layer = [block_stride]
+        self.num_regions = 1
+        self.num_descs = self.num_blocks
+        self.kv_caches_base_addr[self.engine_id][self.tp_rank] = [base_addr]
+
+        descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)
+        self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
+        self._registered_descs.append(descs)
+
+        self.dst_num_blocks[self.engine_id] = self.num_blocks
+
+        self.src_xfer_handles_by_block_size[self.block_size], (self.src_blocks_data) = (
+            self.register_local_xfer_handler(self.block_size)
+        )
+
+        agent_metadata = NixlAgentMetadata(
+            engine_id=self.engine_id,
+            agent_metadata=self.nixl_wrapper.get_agent_metadata(),
+            device_id=self.device_id,
+            kv_caches_base_addr=(
+                self.kv_caches_base_addr[self.engine_id][self.tp_rank]
+            ),
+            num_blocks=self.num_blocks,
+            block_lens=self.block_len_per_layer,
+            kv_cache_layout=self.kv_cache_layout,
+            block_size=self.block_size,
+            ssm_sizes=self._mamba_ssm_size,
+            attn_backend_name=self.backend_name,
+            physical_blocks_per_logical_kv_block=(
+                self._physical_blocks_per_logical_kv_block
+            ),
+        )
+        assert self.compat_hash is not None
+        encoder = msgspec.msgpack.Encoder()
+        self.xfer_handshake_metadata = NixlHandshakePayload(
+            compatibility_hash=self.compat_hash,
+            agent_metadata_bytes=encoder.encode(agent_metadata),
+        )
+
    def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
        """Register the KV Cache data in nixl."""
+
+        # Detect packed allocation: all tensors are strided views into the
+        # same backing storage (different data_ptr but same storage).
+        # This happens with DSv4-style contiguous per-block packing.
+        if len(kv_caches) > 1 and not self._has_mamba:
+            storage = next(iter(kv_caches.values())).untyped_storage()
+            storage_ptrs = {
+                cache.untyped_storage().data_ptr() for cache in kv_caches.values()
+            }
+            data_ptrs = {cache.data_ptr() for cache in kv_caches.values()}
+            if len(storage_ptrs) == 1 and len(data_ptrs) > 1:
+                self._register_packed_kv_cache(storage)
+                self.device_kv_caches = kv_caches
+                return
+
        self.transfer_topo = TransferTopology(
            tp_rank=self.tp_rank,
            tp_size=self.world_size,
@@ -112,7 +112,7 @@ class _StatsKey:

    # Maps metric name -> _MetricType value
    TYPES = "types"
-    # Maps metric name -> observed value (number or list)
+    # Maps metric name -> {label values tuple -> observed value (number or list)}
    DATA = "data"


@@ -125,15 +125,17 @@ class OffloadingConnectorStats(KVConnectorStats):

        {
            _StatsKey.TYPES: {name: _MetricType.*, ...},
-            _StatsKey.DATA:  {name: value, ...},
+            _StatsKey.DATA:  {name: {labelvalues: value, ...}, ...},
        }

    This structure is self-describing: it survives IPC serialization
    without needing the full ``OffloadingMetricMetadata`` objects on the
    receiving side.

-    Counter values are aggregated by summing, gauge values use the latest
-    snapshot, and histogram values are lists of observed samples.
+    Counter values are aggregated by summing per-label-tuple, gauge values
+    use the latest snapshot per-label-tuple, and histogram values are lists of
+    observed samples per-label-tuple. Unlabeled metrics use ``()`` as their
+    labelvalues tuple.
    """

    def __post_init__(self):
@@ -160,26 +162,32 @@ class OffloadingConnectorStats(KVConnectorStats):
        assert isinstance(other, OffloadingConnectorStats)
        other_types = other._types
        other_values = other._values
-        for key, value in other_values.items():
+        for key, other_label_values in other_values.items():
            type_str = other_types.get(key)
            if type_str is None:
                raise AssertionError(f"Unknown offloading stats key: {key}")
            self._types.setdefault(key, type_str)
-            if type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                if key not in self._values:
-                    self._values[key] = value
+            current_label_values = self._values.setdefault(key, {})
+            for labelvalues, value in other_label_values.items():
+                if type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    if labelvalues not in current_label_values:
+                        current_label_values[labelvalues] = list(value)
+                    else:
+                        assert isinstance(current_label_values[labelvalues], list)
+                        current_label_values[labelvalues].extend(value)
+                elif type_str == _MetricType.COUNTER:
+                    assert isinstance(value, int | float)
+                    current_label_values[labelvalues] = (
+                        current_label_values.get(labelvalues, 0) + value
+                    )
+                elif type_str == _MetricType.GAUGE:
+                    assert isinstance(value, int | float)
+                    current_label_values[labelvalues] = value
                else:
-                    assert isinstance(self._values[key], list)
-                    self._values[key].extend(value)
-            elif type_str == _MetricType.COUNTER:
-                assert isinstance(value, int | float)
-                self._values[key] = self._values.get(key, 0) + value
-            elif type_str == _MetricType.GAUGE:
-                assert isinstance(value, int | float)
-                self._values[key] = value
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
        return self

    def reduce(self) -> dict[str, int | float]:
@@ -190,44 +198,62 @@ class OffloadingConnectorStats(KVConnectorStats):
        stats for the last time interval.
        """
        return_dict: dict[str, int | float] = {}
-        for key, value in self._values.items():
+        for key, label_value_map in self._values.items():
            type_str = self._types.get(key)
            if type_str is None:
                raise AssertionError(f"Unknown offloading stats key: {key}")
-            if type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                return_dict[f"{key}_count"] = len(value)
-                return_dict[f"{key}_sum"] = sum(value)
-            elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
-                assert isinstance(value, int | float)
-                return_dict[key] = value
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+            for labelvalues, value in label_value_map.items():
+                key_with_labels = f"{key}:{labelvalues}" if labelvalues else key
+                if type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    return_dict[f"{key_with_labels}_count"] = len(value)
+                    return_dict[f"{key_with_labels}_sum"] = sum(value)
+                elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
+                    assert isinstance(value, int | float)
+                    return_dict[key_with_labels] = value
+                else:
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
        return return_dict

    def is_empty(self) -> bool:
        return not self.data.get(_StatsKey.DATA)

    def increase_counter(
-        self, counter_name: str, counter_increase_value: int | float
+        self,
+        counter_name: str,
+        counter_increase_value: int | float,
+        labelvalues: tuple[str, ...] = (),
    ) -> None:
        """Increase a counter on the stats payload."""
        self._types.setdefault(counter_name, _MetricType.COUNTER)
-        self._values[counter_name] = (
-            self._values.get(counter_name, 0) + counter_increase_value
+        counter_values = self._values.setdefault(counter_name, {})
+        counter_values[labelvalues] = (
+            counter_values.get(labelvalues, 0) + counter_increase_value
        )

-    def set_gauge(self, gauge_name: str, gauge_value: int | float) -> None:
+    def set_gauge(
+        self,
+        gauge_name: str,
+        gauge_value: int | float,
+        labelvalues: tuple[str, ...] = (),
+    ) -> None:
        """Set a gauge snapshot on the stats payload."""
        self._types.setdefault(gauge_name, _MetricType.GAUGE)
-        self._values[gauge_name] = gauge_value
+        gauge_values = self._values.setdefault(gauge_name, {})
+        gauge_values[labelvalues] = gauge_value

    def observe_histogram(
-        self, histogram_name: str, histogram_value: int | float
+        self,
+        histogram_name: str,
+        histogram_value: int | float,
+        labelvalues: tuple[str, ...] = (),
    ) -> None:
        """Record a histogram observation on the stats payload."""
        self._types.setdefault(histogram_name, _MetricType.HISTOGRAM)
-        self._values.setdefault(histogram_name, []).append(histogram_value)
+        histogram_values = self._values.setdefault(histogram_name, {})
+        histogram_values.setdefault(labelvalues, []).append(histogram_value)


 class OffloadPromMetrics(KVConnectorPromMetrics):
@@ -255,7 +281,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics):

        self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec)
        self._offloading_metric_defs: dict[str, PromMetricT] = {}
-        self.offloading_metrics: dict[tuple[int, str], PromMetricT] = {}
+        # (engine_idx, metric_name, labelvalues) -> metric with bound labels
+        self.offloading_metrics: dict[
+            tuple[int, str, tuple[str, ...]], PromMetricT
+        ] = {}

        self._counter_kv_bytes = self._counter_cls(
            name=_DEPRECATED_TOTAL_BYTES,
@@ -301,10 +330,6 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
            self._offloading_metric_defs[metric_name] = self._create_metric(
                metric_name, metadata
            )
-            for engine_idx, labelvalues in per_engine_labelvalues.items():
-                self.offloading_metrics[(engine_idx, metric_name)] = (
-                    self._offloading_metric_defs[metric_name].labels(*labelvalues)
-                )

    def _create_metric(
        self, metric_name: str, metadata: OffloadingMetricMetadata
@@ -312,7 +337,7 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
        kwargs: dict[str, Any] = {
            "name": metric_name,
            "documentation": metadata.documentation,
-            "labelnames": self._labelnames,
+            "labelnames": self._labelnames + list(metadata.labelnames),
        }
        if isinstance(metadata, OffloadingCounterMetadata):
            metric_cls = self._counter_cls
@@ -326,11 +351,37 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
            raise AssertionError(f"Unknown offloading metric metadata: {metadata}")
        return metric_cls(**kwargs)

+    def _get_prometheus_metric(
+        self,
+        metric_name: str,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
+    ) -> PromMetric:
+        metadata = self._offloading_metric_metadata[metric_name]
+        if len(labelvalues) != len(metadata.labelnames):
+            raise AssertionError(
+                f"Metric {metric_name} expects {len(metadata.labelnames)} labels, "
+                f"got {len(labelvalues)}"
+            )
+        key = (engine_idx, metric_name, labelvalues)
+        prom_metric = self.offloading_metrics.get(key)
+        if prom_metric is None:
+            engine_labelvalues = self.per_engine_labelvalues[engine_idx]
+            prom_metric = self._offloading_metric_defs[metric_name].labels(
+                *(engine_labelvalues + list(labelvalues))
+            )
+            self.offloading_metrics[key] = prom_metric
+        return prom_metric
+
    def _increase_counter(
-        self, metric_name: str, value: int | float, engine_idx: int
+        self,
+        metric_name: str,
+        value: int | float,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
    ) -> None:
-        self.offloading_metrics[(engine_idx, metric_name)].inc(value)
-        if not self._observe_deprecated_metrics:
+        self._get_prometheus_metric(metric_name, labelvalues, engine_idx).inc(value)
+        if labelvalues or not self._observe_deprecated_metrics:
            return
        # Keep deprecated CPU offload transfer metrics updated during the
        # transition to flat metric names.
@@ -343,15 +394,26 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
        elif metric_name == _TransferMetricName.STORE_TIME:
            self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value)

-    def _set_gauge(self, metric_name: str, value: int | float, engine_idx: int) -> None:
-        self.offloading_metrics[(engine_idx, metric_name)].set(value)
+    def _set_gauge(
+        self,
+        metric_name: str,
+        value: int | float,
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
+    ) -> None:
+        self._get_prometheus_metric(metric_name, labelvalues, engine_idx).set(value)

    def _observe_histogram(
-        self, metric_name: str, value: list[int | float], engine_idx: int
+        self,
+        metric_name: str,
+        value: list[int | float],
+        labelvalues: tuple[str, ...],
+        engine_idx: int,
    ) -> None:
+        prom_metric = self._get_prometheus_metric(metric_name, labelvalues, engine_idx)
        for observation in value:
-            self.offloading_metrics[(engine_idx, metric_name)].observe(observation)
-            if not self._observe_deprecated_metrics:
+            prom_metric.observe(observation)
+            if labelvalues or not self._observe_deprecated_metrics:
                continue
            # Keep deprecated CPU offload transfer metrics updated during the
            # transition to flat metric names.
@@ -368,20 +430,23 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
        """Observe transfer statistics."""
        metric_types = transfer_stats_data.get(_StatsKey.TYPES, {})
        metric_data = transfer_stats_data.get(_StatsKey.DATA, {})
-        for key, value in metric_data.items():
+        for key, label_value_map in metric_data.items():
            type_str = metric_types.get(key)
            if type_str is None:
                raise AssertionError(f"Unknown offloading stats key: {key}")
            assert key in self._offloading_metric_defs
-            if type_str == _MetricType.COUNTER:
-                assert isinstance(value, int | float)
-                self._increase_counter(key, value, engine_idx)
-            elif type_str == _MetricType.GAUGE:
-                assert isinstance(value, int | float)
-                self._set_gauge(key, value, engine_idx)
-            elif type_str == _MetricType.HISTOGRAM:
-                assert isinstance(value, list)
-                assert all(isinstance(v, int | float) for v in value)
-                self._observe_histogram(key, value, engine_idx)
-            else:
-                raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
+            for labelvalues, value in label_value_map.items():
+                if type_str == _MetricType.COUNTER:
+                    assert isinstance(value, int | float)
+                    self._increase_counter(key, value, labelvalues, engine_idx)
+                elif type_str == _MetricType.GAUGE:
+                    assert isinstance(value, int | float)
+                    self._set_gauge(key, value, labelvalues, engine_idx)
+                elif type_str == _MetricType.HISTOGRAM:
+                    assert isinstance(value, list)
+                    assert all(isinstance(v, int | float) for v in value)
+                    self._observe_histogram(key, value, labelvalues, engine_idx)
+                else:
+                    raise AssertionError(
+                        f"Unknown metric type '{type_str}' for key: {key}"
+                    )
@@ -647,6 +647,13 @@ class OffloadingConnectorScheduler:
        for group_state in req_status.group_states:
            group_state.block_ids.clear()

+        if req_status.transfer_jobs:
+            logger.debug(
+                "Delaying request %s since it still has in-flight transfers",
+                request.request_id,
+            )
+            return None, False
+
        req_status.update_offload_keys()
        req_status.num_locally_computed_tokens = num_computed_tokens

@@ -50,7 +50,8 @@ class OffloadingConnectorWorker:
    def register_kv_caches(
        self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]
    ):
-        num_blocks = self.spec.kv_cache_config.num_blocks
+        kv_cache_config = self.spec.kv_cache_config
+        num_blocks = kv_cache_config.num_blocks

        # layer_name -> (num_blocks, page_size_bytes) tensor
        tensors_per_block: dict[str, tuple[torch.Tensor, ...]] = {}
@@ -58,7 +59,7 @@ class OffloadingConnectorWorker:
        unpadded_page_size_bytes: dict[str, int] = {}
        # layer_name -> size of page in bytes
        page_size_bytes: dict[str, int] = {}
-        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
            group_layer_names = kv_cache_group.layer_names
            group_kv_cache_spec = kv_cache_group.kv_cache_spec
            if isinstance(group_kv_cache_spec, UniformTypeKVCacheSpecs):
@@ -72,18 +73,22 @@ class OffloadingConnectorWorker:
                if isinstance(layer_kv_cache_spec, AttentionSpec):
                    layer_kv_cache = kv_caches[layer_name]
                    assert isinstance(layer_kv_cache, torch.Tensor)
-                    assert layer_kv_cache.storage_offset() == 0

-                    storage = layer_kv_cache.untyped_storage()
                    page = layer_kv_cache_spec.page_size_bytes
+                    elem_size = layer_kv_cache.element_size()
+                    byte_offset = layer_kv_cache.storage_offset() * elem_size
+                    block_stride_bytes = layer_kv_cache.stride(0) * elem_size
                    tensors_per_block[layer_name] = (
                        torch.tensor(
                            [],
                            dtype=torch.int8,
                            device=layer_kv_cache.device,
-                        )
-                        .set_(storage)
-                        .view(num_blocks, page),
+                        ).set_(
+                            layer_kv_cache.untyped_storage(),
+                            byte_offset,
+                            (num_blocks, page),
+                            (block_stride_bytes, 1),
+                        ),
                    )
                    page_size_bytes[layer_name] = layer_kv_cache_spec.page_size_bytes
                    unpadded_page_size_bytes[layer_name] = (
@@ -118,9 +123,35 @@ class OffloadingConnectorWorker:
                else:
                    raise NotImplementedError

+        packed_kv_cache_tensor = next(
+            (t for t in kv_cache_config.kv_cache_tensors if t.block_stride), None
+        )
+        is_dsv4 = all(
+            isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs)
+            for group in kv_cache_config.kv_cache_groups
+        )
+        if packed_kv_cache_tensor is not None and not is_dsv4:
+            (tensor,) = tensors_per_block[packed_kv_cache_tensor.shared_by[0]]
+            block_stride = tensor.stride(0)
+            packed_tensor = tensor.as_strided(
+                (num_blocks, block_stride),
+                (block_stride, 1),
+                storage_offset=0,
+            )
+            self._register_handlers(
+                CanonicalKVCaches(
+                    [CanonicalKVCacheTensor(packed_tensor, block_stride)],
+                    [
+                        [CanonicalKVCacheRef(0, block_stride)]
+                        for _ in kv_cache_config.kv_cache_groups
+                    ],
+                )
+            )
+            return
+
        block_tensors: list[CanonicalKVCacheTensor] = []
        block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)
-        for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors:
+        for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
            # Filter to layers that were actually processed above.
            # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for
            # every (tuple_idx, page_size) slot; slots where no group has a
@@ -162,7 +193,7 @@ class OffloadingConnectorWorker:
                    )

        group_data_refs: list[list[CanonicalKVCacheRef]] = []
-        for kv_cache_group in self.spec.kv_cache_config.kv_cache_groups:
+        for kv_cache_group in kv_cache_config.kv_cache_groups:
            group_refs: list[CanonicalKVCacheRef] = []
            for layer_name in kv_cache_group.layer_names:
                group_refs += block_data_refs[layer_name]
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Woosuk Kwon	fb9f5790bb	megamoe Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 23:59:53 +00:00
Woosuk Kwon	57b8526cfb	Generalize Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 19:19:36 +00:00
Woosuk Kwon	25411f3138	minor Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 04:57:51 +00:00
Woosuk Kwon	261a8820c7	Bound DeepEPV2 num_max_tokens_per_rank Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 04:52:12 +00:00
Woosuk Kwon	416977534a	fuse expert id gather Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 04:51:16 +00:00
Woosuk Kwon	ca54c027c4	minor Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 01:51:39 +00:00
Woosuk Kwon	65e05df079	minor Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 01:34:59 +00:00
Woosuk Kwon	2a09d50034	Merge branch 'main' into woosuk/triton-fix	2026-06-22 00:29:09 +00:00
Woosuk Kwon	6f3d89d105	masked moe sum Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-22 00:26:33 +00:00
MattandGitHub	a19ff2218a	[Hardware][AMD][CI] Fix Spec Decode Eagle test group (#46018 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-06-21 17:40:02 -05:00
MattandGitHub	4f0d0049a0	[Hardware][AMD][CI] Fix Kernels Attention test groups (#46080 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-06-21 17:10:51 -05:00
Charlie Fu GitHub Andreas Karatzas	13b83d77ad	[ROCm][CI] skip test_double_aiter_rms_quant_fusion (#45967 ) Signed-off-by: charlifu <charlifu@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>	2026-06-21 16:53:11 -05:00
MattandGitHub	50241602fd	[Hardware][AMD][CI] Fix gfx942 Kernels MoE test group (#46298 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-06-21 16:45:37 -05:00
Ting SUNandGitHub	12fe2a9aac	[Bugfix][Qwen3-VL] Fix multi-video crash with list-valued fps/num_frames (#46305 ) Signed-off-by: Ting Sun <suntcrick@gmail.com>	2026-06-21 14:31:23 -07:00
Benjamin ChislettandGitHub	89bd2c14d3	[Spec Decode] Add Qwen3 architecture support for EAGLE3 (#43132 ) Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>	2026-06-21 13:55:26 -07:00
ZedongLiuandGitHub	9c450b1027	[Kernel][Bugfix] Fix INT8 per-token-head KV cache rounding in Triton reshape-and-cache (#45361 ) Signed-off-by: ZedongLiu <113341356+Zedong-Liu@users.noreply.github.com>	2026-06-21 15:59:40 -04:00
Ranran GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Isotr0py	635c38338a	[Multimodal] Add Qwen2-VL/Qwen2.5-VL processor-mapped video loader (#45555 ) Signed-off-by: Ranran <hzz5361@psu.edu> Signed-off-by: Ranran Haoran Zhang <ranzhang@redhat.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>	2026-06-21 18:56:50 +00:00
Srinivas Krovvidi GitHub srinivas_oo7	c441ad1c07	[KV Offloading] Add labeled metrics support (#45957 ) Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>	2026-06-21 18:04:01 +00:00
Jee Jee LiandGitHub	745bba5ea8	[Model]Fix MiniMaxM2ForCausalLM perf regression (#45935 ) Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>	2026-06-22 00:28:52 +08:00
Alex Steiner GitHub Claude Opus 4.8 Giancarlo Delfin Yifan Qiao	2cac89f9da	[Spec Decode] Support mixed KV page sizes for DFlash (#45181 ) Signed-off-by: Alex Steiner <asteiner@nvidia.com> Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>	2026-06-21 22:45:14 +08:00
aoshen02 GitHub Claude Opus 4.8 Roger Wang	3e6e33526d	[Disagg] return routed_experts on streaming generate responses (#44638 ) Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Roger Wang <hey@rogerw.io>	2026-06-21 07:37:10 -07:00
junkang1991 GitHub Hongxia Yang Tan Pin Siang vllmellm Chun Fang TianDi101 functionstackx tjtanaamergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	b91b7726e0	[ROCm][P/D] Support MiniMax-M3 mixed KV layouts in MoRIIO READ mode (#46039 ) Signed-off-by: Jun Kang Chow <junkangchow@gmail.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Hongxia Yang <hongxia.yang@amd.com> Co-authored-by: Tan Pin Siang <tanpinsiang@gmail.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: Chun Fang <chun.fang@amd.com> Co-authored-by: TianDi101 <ditian12@amd.com> Co-authored-by: functionstackx <47992694+functionstackx@users.noreply.github.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-06-21 12:55:19 +00:00
Palaiologos1453andGitHub	d3ad8e8bcd	[Bugfix] Defer offload reads while transfers are pending (#46231 ) Signed-off-by: test test <2260891073@qq.com>	2026-06-21 14:30:13 +03:00
xiaolinchen GitHub wentian-byte	b80ce9dd2f	[CI][test] Replace InternVL2-1B with InternVL3-1B in test_pipeline_parallel.py (#46241 ) Signed-off-by: wentian-byte <192079369+wentian-byte@users.noreply.github.com> Co-authored-by: wentian-byte <192079369+wentian-byte@users.noreply.github.com>	2026-06-21 15:11:19 +08:00
Shifani Rajabose GitHub Kunshang Ji	b5495cc5f9	Fix memory pointer overflow in Mamba state buffers (#44665 ) Signed-off-by: Shifani Rajabose <shifani.rajabose@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-06-21 14:00:50 +08:00
Ting SUNandGitHub	183a430c13	[Bugfix][Model Runner V2] Fix min_tokens off-by-one in the V2 GPU sampler (#46243 ) Signed-off-by: Ting Sun <suntcrick@gmail.com>	2026-06-21 05:06:49 +00:00
MattandGitHub	a346d589f5	[Bugfix] Fix NVFP4/OCP MX MoE emulation (#46254 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>	2026-06-20 23:13:10 -05:00
Nick HillandGitHub	7df3d7dada	[Core] Ensure memory is pinned prior to async h2d copy (#45424 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-06-20 20:02:24 -07:00
Umut Polat GitHub Flora Feng	8dd1b702f2	[Misc] Fix stale doc URL and docstring module path (#35530 ) Signed-off-by: umut-polat <52835619+umut-polat@users.noreply.github.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>	2026-06-20 23:57:01 +00:00
aoshen02 GitHub Martin Hickey	f57ac274b2	[Render] Add reasoning/tool parsing to /derender + fix byte-fallback FFFD (#45919 ) Signed-off-by: aoshen524 <aoshen524@gmail.com> Co-authored-by: Martin Hickey <martin.hickey@ie.ibm.com>	2026-06-20 19:43:32 -04:00
aman GitHub Claude Nick Hill	6e919960af	[Perf] Skip/shrink all_token_ids copy in scheduler for non-async and V2 runner (#45840 ) Signed-off-by: amanchugh89 <amanchugh.89@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>	2026-06-20 22:36:57 +00:00
Jonathan ChenandGitHub	c88d3d4775	[SimpleCPUOffloadConnector] PCP + DCP support (#39831 ) Signed-off-by: Jonathan Chen <chenleejonathan@gmail.com>	2026-06-20 15:01:06 -07:00
Yifan QiaoandGitHub	ab7fcbdd5d	[Perf][KVConnector][Mooncake] Compact chunk-hash keys and zero-copy lookup wire format (#45969 )	2026-06-20 15:00:11 -07:00
Varun Sundar Rabindranath GitHub Varun Sundar Rabindranath	3b4a76b63f	[KV-Offloading] : Expose CPU cache usage metric (#45737 ) Signed-off-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com> Signed-off-by: <> Co-authored-by: Varun Sundar Rabindranath <varun-sundar-rabindranath@h100-01.nemg-001.lab.rdu2.dc.redhat.com>	2026-06-20 21:21:55 +00:00
Lucas Wilkinson GitHub OpenAI Codex Tyler Michael Smith	cc22621b51	[KV Offload] Support packed HMA KV cache layout (#46205 ) Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-06-20 21:19:40 +00:00
Tyler Michael Smith GitHub Claude Opus 4.6	77148992cf	[Bugfix] Move extract_layer_index back inside is_v32 guard (#46199 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>	2026-06-20 21:19:10 +00:00
shuoming zhang GitHub Tyler Michael Smith	891cc4b9c5	[Frontend] Report cache usage in Anthropic /v1/messages API (#40912 ) Signed-off-by: mistral0105 <zhangshuoming17@mails.ucas.ac.cn> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-06-20 21:12:48 +00:00
TJianandGitHub	1bdf9810aa	[ROCm] [Bugfix] Bugfix ROCm Sparse Indexer (#46222 ) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>	2026-06-20 13:38:42 -07:00
Tyler Michael Smith GitHub Claude Codexmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>kourosh hakhamaneshi	ebfbcfe46a	Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: kourosh hakhamaneshi <kouroshHakha@users.noreply.github.com>	2026-06-20 13:38:10 -07:00
Tyler Michael Smith GitHub Claude Opus 4.6	e9de72fe6c	[Bugfix] Guard model_config access in _log_compilation_config (#46198 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>	2026-06-20 19:26:38 +00:00
L丶 GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	d272418f45	[Perf] Optimize Qwen3-VL multi-video prompt processing (#46026 ) Signed-off-by: Sirius29 <422058530@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-06-20 07:09:18 -07:00
Sumanth R HegdeandGitHub	7ff7f5c8eb	Revert "Fix Stale Encoder Cache After Weight Update" (#46125 )	2026-06-20 07:09:09 -07:00
Matt GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	dced290769	[Hardware][AMD][CI] Fix e2e core test group (#46024 ) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-06-20 02:04:35 -05:00
JasonLi314andGitHub	93bad11912	[Bugfix] Fix gridDim.y overflow for large row counts (#45255 ) Signed-off-by: Jason Li <li.jason.cs@gmail.com>	2026-06-19 23:27:45 -04:00
djramicandGitHub	0fbf42af84	[ROCm] Fix VRAM not freed in test_phi3v (#46046 ) Signed-off-by: Djordje Ramic <djoramic@amd.com>	2026-06-19 17:20:59 -05:00
Charlie FuandGitHub	e6cd8913dd	[ROCm][CI] Skip Qwen3.5-35B-A3B-MXFP4-AITER-TP2 for non gfx950 (#46109 ) Signed-off-by: charlifu <charlifu@amd.com>	2026-06-19 17:20:10 -05:00
Ben BrowningandGitHub	859e4d436b	[Bugfix][Parser] Fix U+FFFD leak at reasoning-to-content transition in engine parsers (#46159 ) Signed-off-by: Ben Browning <bbrownin@redhat.com>	2026-06-19 22:09:28 +00:00
Micah WilliamsonandGitHub	4a083cc858	[ROCm][CI] Pin `test_rocm_compressed_tensors_w8a8` to TRITON_ATTN (#46180 ) Signed-off-by: Micah Williamson <micah.williamson@amd.com>	2026-06-19 15:20:06 -05:00
Vadim GimpelsonandGitHub	ca7e1f2c43	Move CI failure diagnosis docs into ci-fails-buildkite skill (#45975 ) Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>	2026-06-19 20:12:40 +00:00
djramicandGitHub	dec860fb19	[ROCm] Use vLLM's fp8 quant max in AITER hipBLASLt accuracy test (#46176 ) Signed-off-by: Djordje Ramic <djoramic@amd.com>	2026-06-19 13:24:02 -05:00
Harry MellorandGitHub	0a49fb2b13	Fix dead link in docs (#46181 ) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>	2026-06-19 18:16:09 +00:00
Ben BrowningandGitHub	4a8abf37c7	[Test] Migrate test_openai_schema.py to schemathesis 4.x (#46173 ) Signed-off-by: Ben Browning <bbrownin@redhat.com>	2026-06-19 18:05:18 +00:00
Tyler Michael Smith GitHub Claude Opus 4.6 Matthew Bonanni Lucas Wilkinson Lucas Wilkinson OpenAI Codex	01192139bf	[DSv4] Pack KV caches into contiguous per-block allocations for DeepSeek V4 (#44577 ) Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-06-19 12:55:42 -04:00
Chris LeonardandGitHub	b9a7cd464c	[12/n] final _C library kernel migration (#45415 )	2026-06-19 06:57:26 -07:00
Woosuk Kwon	a06a16ff0a	minor Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-19 03:31:41 +00:00
Woosuk Kwon	a1d80989d9	Plumb is_padding Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-19 03:23:23 +00:00
Woosuk Kwon	40a19bed77	Merge branch 'main' into woosuk/triton-fix	2026-06-18 16:06:08 +00:00
Woosuk Kwon	8e8f8c2e2c	remove batched Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-18 01:16:13 +00:00
Woosuk Kwon	116d3e3149	Merge branch 'main' into woosuk/triton-fix	2026-06-18 01:14:50 +00:00
Woosuk Kwon	f6fa9700e6	Merge branch 'main' into woosuk/triton-fix	2026-06-17 23:38:48 +00:00
Woosuk Kwon	0221ab433e	wip Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>	2026-06-17 22:57:43 +00:00