Merge branch 'main' into cuda-arch-fixup

Update qutlass cmake for stable abi (#47879 )
Signed-off-by: Chris Leonard <chleonar@redhat.com>
2026-07-21 09:33:48 +08:00 · 2026-07-20 18:31:10 -07:00 · 2026-07-20 18:30:55 -07:00 · 2026-07-20 18:30:38 -07:00 · 2026-07-21 01:05:05 +00:00 · 2026-07-21 00:29:10 +00:00
168 changed files with 7165 additions and 1668 deletions
@@ -18,6 +18,8 @@ steps:
  - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
  - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
  - tests/kernels/mamba/test_cpu_short_conv.py
+  - tests/kernels/mamba/test_causal_conv1d.py
+  - tests/kernels/mamba/test_mamba_ssm.py
  commands:
    - |
      bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -28,7 +30,9 @@ steps:
      pytest -x -v -s tests/kernels/test_onednn.py
      pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
      pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
-      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
+      pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+      pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
+      pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"

 # Note: SDE can't be downloaded from CI host because of AWS WAF
 # - label: CPU-Compatibility Tests
@@ -813,8 +813,8 @@ steps:

            # Download artifacts from current build
            echo "Downloading artifacts from current build"
-            # buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
-            # buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
+            buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
+            buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .

            # # Run upload script
            bash .buildkite/scripts/upload-rocm-wheels.sh
@@ -40,7 +40,9 @@ function cpu_tests() {
    pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
    pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
    pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
-    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
+    pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
+    pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
+    pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"

  # skip tests requiring model downloads if HF_TOKEN is not set
  # due to rate-limits
@@ -97,3 +99,4 @@ function cpu_tests() {
 # All of CPU tests are expected to be finished less than 40 mins.
 export -f cpu_tests
 timeout 2h bash -c cpu_tests
+
@@ -10,7 +10,9 @@ steps:
  - vllm/engine/arg_utils.py
  - vllm/config/model.py
  - vllm/model_executor
+  - vllm/model_executor/warmup
  - tests/model_executor
+  - tests/model_executor/test_jit_warmup.py
  - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
  commands:
    - apt-get update && apt-get install -y curl libsodium23
@@ -34,7 +36,9 @@ steps:
      - vllm/engine/arg_utils.py
      - vllm/config/model.py
      - vllm/model_executor
+      - vllm/model_executor/warmup
      - tests/model_executor
+      - tests/model_executor/test_jit_warmup.py
      - tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
      - vllm/_aiter_ops.py
      - vllm/platforms/rocm.py
@@ -227,6 +227,22 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
  cuda_archs_loose_intersection(CUDA_ARCHS
    "${CUDA_SUPPORTED_ARCHS}" "${CUDA_ARCHS}")
  message(STATUS "CUDA supported target architectures: ${CUDA_ARCHS}")
+
+  set(VLLM_COMPILED_CUDA_ARCHS)
+  foreach(_ARCH ${CUDA_ARCHS})
+    set(_COMPILED_ARCH "${_ARCH}")
+    if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
+      if(_ARCH MATCHES "^(10|11|12)\\.0$")
+        set(_COMPILED_ARCH "${_ARCH}f")
+      endif()
+    elseif(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
+      if(_ARCH MATCHES "^(10\\.(0|1|3)|12\\.(0|1))$")
+        set(_COMPILED_ARCH "${_ARCH}a")
+      endif()
+    endif()
+    list(APPEND VLLM_COMPILED_CUDA_ARCHS "${_COMPILED_ARCH}")
+  endforeach()
+  list(JOIN VLLM_COMPILED_CUDA_ARCHS "," VLLM_COMPILED_CUDA_ARCHS_STR)
 else()
  #
  # For other GPU targets override the GPU architectures detected by cmake/torch
@@ -385,8 +401,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  #
  # _C_stable_libtorch extension (ops registered via STABLE_TORCH_LIBRARY)
  #
+  # Shared entry sources are part of the base extension source list, but some
+  # optional kernel families below append source-local feature macros to them.
+  set(VLLM_STABLE_TORCH_BINDINGS_SRC
+    "csrc/libtorch_stable/torch_bindings.cpp")
+  set(CACHE_KERNELS_SRC "csrc/libtorch_stable/cache_kernels.cu")
+  set(SCALED_MM_ENTRY_SRC
+    "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu")
+  set(NVFP4_QUANT_ENTRY_SRC
+    "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu")
+  set(NVFP4_SCALED_MM_ENTRY_SRC
+    "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu")
+
  set(VLLM_STABLE_EXT_SRC
-    "csrc/libtorch_stable/torch_bindings.cpp"
+    "${VLLM_STABLE_TORCH_BINDINGS_SRC}"
    "csrc/libtorch_stable/cuda_view.cu"
    "csrc/libtorch_stable/cuda_utils_kernels.cu"
    "csrc/libtorch_stable/activation_kernels.cu"
@@ -409,7 +437,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    "csrc/libtorch_stable/sampler.cu"
    "csrc/libtorch_stable/topk.cu"
    "csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
-    "csrc/libtorch_stable/cache_kernels.cu"
+    "${CACHE_KERNELS_SRC}"
    "csrc/libtorch_stable/cache_kernels_fused.cu"
    "csrc/libtorch_stable/custom_all_reduce.cu"
    "csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
@@ -425,11 +453,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      cuda_archs_loose_intersection(COOPERATIVE_TOPK_ARCHS
        "9.0a;10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
    endif()
-
-    if(COOPERATIVE_TOPK_ARCHS)
-      list(APPEND VLLM_GPU_FLAGS "-DVLLM_ENABLE_COOPERATIVE_TOPK=1")
-
-    endif()
  endif()

  if(VLLM_GPU_LANG STREQUAL "CUDA")
@@ -467,11 +490,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

    list(APPEND VLLM_STABLE_EXT_SRC
      "csrc/libtorch_stable/cutlass_extensions/common.cpp"
-      "csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_entry.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"
+      "${SCALED_MM_ENTRY_SRC}"
+      "${NVFP4_QUANT_ENTRY_SRC}"
+      "${NVFP4_SCALED_MM_ENTRY_SRC}"
      "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu"
      "csrc/libtorch_stable/minimax_reduce_rms_kernel.cu")
+    set_compile_definitions_for_srcs(
+      SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC}"
+      DEFINITIONS VLLM_COMPILED_CUDA_ARCHS=\"${VLLM_COMPILED_CUDA_ARCHS_STR}\")

    #
    # Machete kernels
@@ -547,11 +573,15 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      CUDA_ARCHS "${CUDA_ARCHS}")

    if(COOPERATIVE_TOPK_ARCHS)
+      set(COOPERATIVE_TOPK_SRC "csrc/libtorch_stable/cooperative_topk.cu")
      list(APPEND VLLM_STABLE_EXT_SRC
-        "csrc/libtorch_stable/cooperative_topk.cu")
+        "${COOPERATIVE_TOPK_SRC}")
      set_gencode_flags_for_srcs(
-        SRCS "csrc/libtorch_stable/cooperative_topk.cu"
+        SRCS "${COOPERATIVE_TOPK_SRC}"
        CUDA_ARCHS "${COOPERATIVE_TOPK_ARCHS}")
+      set_compile_definitions_for_srcs(
+        SRCS "${VLLM_STABLE_TORCH_BINDINGS_SRC};${COOPERATIVE_TOPK_SRC}"
+        DEFINITIONS VLLM_ENABLE_COOPERATIVE_TOPK=1)
    endif()

    # Only build Marlin kernels if we are building for at least some compatible archs.
@@ -760,8 +790,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set_gencode_flags_for_srcs(
      SRCS "${SCALED_MM_SM90_SRCS}"
      CUDA_ARCHS "${SCALED_MM_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM90_SRCS}"
+      DEFINITIONS ENABLE_SCALED_MM_SM90=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM90_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM90=1")
    # Let scaled_mm_c2x know it doesn't need to build these arches
    list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
    message(STATUS "Building scaled_mm_c3x_sm90 for archs: ${SCALED_MM_ARCHS}")
@@ -794,8 +826,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set_gencode_flags_for_srcs(
      SRCS "${SCALED_MM_SM120_SRCS}"
      CUDA_ARCHS "${SCALED_MM_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM120_SRCS}"
+      DEFINITIONS ENABLE_SCALED_MM_SM120=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM120_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM120=1")
    # Let scaled_mm_c2x know it doesn't need to build these arches
    list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
    message(STATUS "Building scaled_mm_c3x_sm120 for archs: ${SCALED_MM_ARCHS}")
@@ -828,8 +862,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set_gencode_flags_for_srcs(
      SRCS "${SCALED_MM_SM100_SRCS}"
      CUDA_ARCHS "${SCALED_MM_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_SM100_SRCS}"
+      DEFINITIONS ENABLE_SCALED_MM_SM100=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_SM100_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_SM100=1")
    # Let scaled_mm_c2x know it doesn't need to build these arches
    list(APPEND SCALED_MM_3X_ARCHS "${SCALED_MM_ARCHS}")
    message(STATUS "Building scaled_mm_c3x_sm100 for archs: ${SCALED_MM_ARCHS}")
@@ -858,8 +894,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set_gencode_flags_for_srcs(
      SRCS "${SCALED_MM_C2X_SRCS}"
      CUDA_ARCHS "${SCALED_MM_2X_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${SCALED_MM_ENTRY_SRC};${SCALED_MM_C2X_SRCS}"
+      DEFINITIONS ENABLE_SCALED_MM_C2X=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${SCALED_MM_C2X_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_SCALED_MM_C2X=1")
    message(STATUS "Building scaled_mm_c2x for archs: ${SCALED_MM_2X_ARCHS}")
  else()
    if (SCALED_MM_3X_ARCHS)
@@ -884,8 +922,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set_gencode_flags_for_srcs(
      SRCS "${CUTLASS_MOE_SM90_SRCS}"
      CUDA_ARCHS "${SCALED_MM_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM90_SRCS}"
+      DEFINITIONS ENABLE_CUTLASS_MOE_SM90=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM90_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
    message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
  else()
    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
@@ -908,8 +948,13 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set_gencode_flags_for_srcs(
      SRCS "${CUTLASS_MOE_SM100_SRCS}"
      CUDA_ARCHS "${SCALED_MM_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${SCALED_MM_ENTRY_SRC};${CUTLASS_MOE_SM100_SRCS}"
+      DEFINITIONS ENABLE_CUTLASS_MOE_SM10X_OR_SM11X=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MOE_SM100_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
+    # The implementation is named sm100 historically, but it is built for the
+    # SM10x/SM11x family: CUDA 12 Thor reports SM101, CUDA 13 Thor reports
+    # SM110. Keep the compile-time macro aligned with runtime dispatch.
    message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
  else()
    if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND SCALED_MM_ARCHS)
@@ -950,10 +995,16 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # FP4/NVFP4 kernels (moved from _C to _C_stable_libtorch)
  #

-  # SM12x FP4 kernels. These share some generic NVFP4 quantization entry
-  # sources with the SM10x/11x block below; set_gencode_flags_for_srcs appends
-  # per-source flags, so shared files accumulate both SM12x and SM10x/11x
-  # gencodes when both families are requested.
+  # Shared FP4 implementation sources live next to the FP4 arch logic because
+  # both SM12x and SM10x/11x append family-specific gencodes and feature macros
+  # to them.
+  set(FP4_SHARED_SRCS
+    "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
+    "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
+    "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
+    "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
+    "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
+
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
    cuda_archs_loose_intersection(FP4_SM120_ARCHS "12.0f" "${CUDA_ARCHS}")
  else()
@@ -961,18 +1012,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  endif()
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM120_ARCHS)
    set(FP4_SM120_SRCS
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
+      ${FP4_SHARED_SRCS}
      "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_sm120_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
-      "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
+    )
    set_gencode_flags_for_srcs(
      SRCS "${FP4_SM120_SRCS}"
      CUDA_ARCHS "${FP4_SM120_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM120_SRCS}"
+      DEFINITIONS ENABLE_NVFP4_SM120=1)
+    set_compile_definitions_for_srcs(
+      SRCS "${SCALED_MM_ENTRY_SRC}"
+      DEFINITIONS ENABLE_CUTLASS_MOE_SM120=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM120_SRCS}")
-    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}")
  else()
    message(STATUS "Not building SM12x NVFP4 as no compatible archs were found.")
@@ -987,14 +1039,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  endif()
  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND FP4_SM100_ARCHS)
    set(FP4_SM100_SRCS
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
+      ${FP4_SHARED_SRCS}
      "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
      "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
-      "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu"
-      "csrc/libtorch_stable/nvfp4_kv_cache_kernels.cu")
+      "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
    if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.9)
      message(STATUS
        "Building mxfp4_experts_quant unsupported stubs because CUDA compiler version is not >= 12.9 (found ${CMAKE_CUDA_COMPILER_VERSION}).")
@@ -1002,9 +1050,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
    set_gencode_flags_for_srcs(
      SRCS "${FP4_SM100_SRCS}"
      CUDA_ARCHS "${FP4_SM100_ARCHS}")
+    set_compile_definitions_for_srcs(
+      SRCS "${NVFP4_QUANT_ENTRY_SRC};${NVFP4_SCALED_MM_ENTRY_SRC};${CACHE_KERNELS_SRC};${FP4_SM100_SRCS}"
+      DEFINITIONS ENABLE_NVFP4_SM100=1)
    list(APPEND VLLM_STABLE_EXT_SRC "${FP4_SM100_SRCS}")
-    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}")
  else()
    message(STATUS "Not building SM10x/11x NVFP4/MXFP4 as no compatible archs were found.")
@@ -1058,7 +1107,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
      SRCS "${CUTLASS_MLA_SRCS}"
      CUDA_ARCHS "${MLA_ARCHS}")
    list(APPEND VLLM_STABLE_EXT_SRC "${CUTLASS_MLA_SRCS}")
-    list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
    # Add MLA-specific include directories only to MLA source files
    set_source_files_properties(${CUTLASS_MLA_SRCS}
      PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
@@ -1106,10 +1154,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
  # Needed to use cuda/hip APIs from C-shim
  if(VLLM_GPU_LANG STREQUAL "CUDA")
    target_compile_definitions(_C_stable_libtorch PRIVATE USE_CUDA)
-    if(COOPERATIVE_TOPK_ARCHS)
-      target_compile_definitions(_C_stable_libtorch PRIVATE
-        VLLM_ENABLE_COOPERATIVE_TOPK=1)
-    endif()
    # Needed by CUTLASS kernels
    target_compile_definitions(_C_stable_libtorch PRIVATE
      CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
    "csrc/cpu/layernorm.cpp"
    "csrc/cpu/mla_decode.cpp"
    "csrc/cpu/pos_encoding.cpp"
+    "csrc/cpu/mamba_cpu.cpp"
    "csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
    "csrc/cpu/cpu_attn.cpp"
    "csrc/cpu/torch_bindings.cpp")
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
+        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
        "csrc/cpu/activation.cpp"
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
        "csrc/cpu/utils.cpp"
        "csrc/cpu/spec_decode_utils.cpp"
        "csrc/cpu/cpu_attn.cpp"
+        "csrc/cpu/mamba_cpu.cpp"
        "csrc/cpu/dnnl_kernels.cpp"
        "csrc/cpu/torch_bindings.cpp"
        # TODO: Remove these files
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
  set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
 else()
  set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
-  set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
+  set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")

  set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
  if(NOT _qutlass_fc_root)
@@ -125,8 +125,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    CUDA_ARCHS "${QUTLASS_ARCHS}"
  )

-  # 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
@@ -139,9 +137,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
    WITH_SOABI)

  target_compile_definitions(_qutlass_C PRIVATE
-    QUTLASS_DISABLE_PYBIND=1
+    QUTLASS_MINIMAL_BUILD=1
    TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
-    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
+    CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
+    TORCH_TARGET_VERSION=0x020B000000000000ULL
+    USE_CUDA)

  set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
    $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
@@ -39,7 +39,7 @@ else()
  FetchContent_Declare(
          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
+          GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
          GIT_PROGRESS TRUE
          # Don't share the vllm-flash-attn build between build types
          BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
@@ -344,6 +344,28 @@ macro(set_gencode_flags_for_srcs)
  endif()
 endmacro()

+#
+# For a list of source files append preprocessor definitions to file-specific
+# compile options. Use this for optional kernel feature macros so toggling one
+# kernel family does not perturb the compile command for every source in the
+# extension target.
+#
+macro(set_compile_definitions_for_srcs)
+  set(options)
+  set(oneValueArgs)
+  set(multiValueArgs SRCS DEFINITIONS)
+  cmake_parse_arguments(arg "${options}" "${oneValueArgs}"
+                        "${multiValueArgs}" ${ARGN})
+
+  foreach(_DEF ${arg_DEFINITIONS})
+    set_property(
+      SOURCE ${arg_SRCS}
+      APPEND PROPERTY
+      COMPILE_DEFINITIONS "${_DEF}"
+    )
+  endforeach()
+endmacro()
+
 #
 # For the given `SRC_CUDA_ARCHS` list of gencode versions in the form
 #  `<major>.<minor>[letter]` compute the "loose intersection" with the
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
    reg.val[1] = fp16_to_fp32_bits(raw_lo);
  }
  float reduce_sum() const {
-    AliasReg ar;
-    ar.reg = reg;
-    float result = 0;
-    unroll_loop<int, VEC_ELEM_NUM>(
-        [&result, &ar](int i) { result += ar.values[i]; });
-
-    return result;
+    // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
+    // Step 1: pairwise sum of the two 4-wide halves
+    __vector float s = vec_add(reg.val[0], reg.val[1]);
+    // Step 2: rotate by 8 bytes (2 floats) and add
+    s = vec_add(s, vec_sld(s, s, 8));
+    // Step 3: rotate by 4 bytes (1 float) and add  => all lanes hold total
+    s = vec_add(s, vec_sld(s, s, 4));
+    return vec_extract(s, 0);
  }
  FP32Vec8 exp() const {
    f32x4x2_t out;
@@ -0,0 +1,285 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
+// mamba_kernels.hpp.
+
+#include "cpu/mamba_kernels.hpp"
+
+#include <ATen/ATen.h>
+#include <torch/library.h>
+#include <c10/util/Optional.h>
+
+#include "cpu_types.hpp"
+
+// ---------------------------------------------------------------------------
+// causal_conv1d_update
+// ---------------------------------------------------------------------------
+at::Tensor causal_conv1d_update_cpu_impl(
+    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
+    const c10::optional<at::Tensor>& bias,
+    const c10::optional<std::string>& activation,
+    const c10::optional<at::Tensor>& conv_state_indices,
+    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
+  bool do_silu = false;
+  if (activation.has_value()) {
+    const std::string& act = activation.value();
+    do_silu = (act == "silu" || act == "swish");
+  }
+
+  at::ScalarType dtype = x.scalar_type();
+
+  // Input x: contiguous in native dtype.
+  at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
+
+  // conv_state: NEVER copy the full paged tensor just for layout reasons.
+  // If the dtype matches we work directly on conv_state (contiguous or not)
+  // by extracting strides and passing them to the kernel.
+  // Only a dtype-conversion copy is made when types differ (rare for BF16).
+  bool state_type_ok = (conv_state.scalar_type() == dtype);
+  at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
+  // state_c and conv_state may be non-contiguous — that is intentional.
+
+  // Weight: coerce to same dtype if needed (should match in practice)
+  at::Tensor w_c =
+      (weight.scalar_type() != dtype)
+          ? weight.to(dtype).contiguous()
+          : (weight.is_contiguous() ? weight : weight.contiguous());
+
+  // Bias stays float32 (small scalar, used only for fp32 accumulation)
+  at::Tensor bias_f32;
+  if (bias.has_value() && bias.value().defined())
+    bias_f32 = bias.value().to(at::kFloat).contiguous();
+
+  int64_t batch = x_c.size(0);
+  int64_t dim = x_c.size(1);
+  int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
+  int64_t width = w_c.size(1);
+  int64_t state_len = state_c.size(2);
+
+  // Extract strides — works for contiguous AND non-contiguous (transposed)
+  // state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
+  // contiguous) stride(1): between conv channels (dim stride) stride(2):
+  // between state elements (=1 when contiguous, =dim when transposed)
+  int64_t stride_s_slot = state_c.stride(0);
+  int64_t stride_s_dim = state_c.stride(1);
+  int64_t stride_s_state = state_c.stride(2);
+
+  at::Tensor out = x_c.clone();  // native dtype, no float32 alloc
+
+  const int32_t* cache_idx_ptr = nullptr;
+  at::Tensor cache_idx_int;
+  if (conv_state_indices.has_value()) {
+    cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
+    cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
+  }
+
+  VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
+    mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
+        x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
+        stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
+        bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
+        out.data_ptr<scalar_t>(), cache_idx_ptr,
+        static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
+        do_silu);
+  });
+
+  // Write back only when a type-conversion copy was made.
+  // Layout-only non-contiguity is handled via strides above — no copy needed.
+  if (!state_type_ok) conv_state.copy_(state_c);
+
+  return out;
+}
+
+// ---------------------------------------------------------------------------
+// selective_state_update
+// ---------------------------------------------------------------------------
+void selective_state_update_cpu_impl(
+    at::Tensor& state,    // (nstates, nheads, dim, dstate)
+    const at::Tensor& x,  // (N, nheads, dim)
+    const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
+    const at::Tensor& C, const c10::optional<at::Tensor>& D,
+    const c10::optional<at::Tensor>& z,
+    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
+    const c10::optional<at::Tensor>& state_batch_indices,
+    const c10::optional<at::Tensor>& dst_state_batch_indices,
+    int64_t null_block_id, at::Tensor& out,
+    const c10::optional<at::Tensor>& num_accepted_tokens,
+    const c10::optional<at::Tensor>& cu_seqlens) {
+  at::ScalarType state_type = state.scalar_type();
+  at::ScalarType input_type = x.scalar_type();
+
+  // x, B, C must be contiguous and match input_type
+  auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+  at::Tensor x_in = ensure_input(x);
+  at::Tensor B_in = ensure_input(B);
+  at::Tensor C_in = ensure_input(C);
+  at::Tensor z_in;
+  if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
+
+  // A, D, dt_bias are float32 model parameters that arrive here as expanded
+  // tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
+  // We need just the scalar value per head as a (nheads,) 1-D array so that
+  // A_ptr[h] in the kernel correctly reads head h's value.
+  //
+  // Strategy: peel trailing expanded (stride=0) dims via .select(), which is
+  // a zero-copy view.  For A: (nheads, head_dim, dstate) strides (1,0,0)
+  //   → .select(2,0) → (nheads, head_dim) strides (1,0)
+  //   → .select(1,0) → (nheads,)          stride  (1,)   ← contiguous, free.
+  // No allocation, no type conversion (A is already float32).
+  auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = t;
+    // Peel trailing dimensions that are broadcast (stride=0 or size=1)
+    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
+    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+
+  at::Tensor A_f32 = to_per_head_1d_f32(A);  // (nheads,) float32
+  at::Tensor D_f32, dt_bias_f32;
+  if (D.has_value() && D.value().defined())
+    D_f32 = to_per_head_1d_f32(D.value());
+  if (dt_bias.has_value() && dt_bias.value().defined())
+    dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
+
+  // dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
+  // the type conversion so we convert head_dim x fewer elements.
+  at::Tensor dt_f32;
+  {
+    // If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
+    // take a zero-copy view of index 0 along that dim first.
+    at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt;  // (N, nheads)
+    at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
+    dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
+  }
+
+  int64_t nheads = state.size(1);
+  int64_t dim = state.size(2);
+  int64_t dstate = state.size(3);
+  int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
+                  ? cu_seqlens.value().size(0) - 1
+                  : x_in.size(0);
+  int64_t ngroups = B_in.size(1);
+
+  // Strides
+  int64_t stride_state_n = state.stride(0);
+  int64_t stride_state_h = state.stride(1);
+  int64_t stride_state_d = state.stride(2);
+  int64_t stride_x_n = x_in.stride(0);
+  int64_t stride_x_h = x_in.stride(1);
+  int64_t stride_dt_n = dt_f32.stride(0);  // dt is (N, nheads)
+  int64_t stride_BC_n = B_in.stride(0);
+  int64_t stride_BC_g = B_in.stride(1);
+  int64_t stride_out_n = out.stride(0);
+  int64_t stride_out_h = out.stride(1);
+
+  // Optional index pointers
+  auto get_int32_ptr =
+      [](const c10::optional<at::Tensor>& opt) -> const int32_t* {
+    return (opt.has_value() && opt.value().defined())
+               ? opt.value().data_ptr<int32_t>()
+               : nullptr;
+  };
+  const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
+  const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
+  const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
+  const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
+
+  // Dispatch on (state_t, input_t, out_t): write directly into `out`
+  // without any intermediate float32 buffer.
+  VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
+    using state_t = scalar_t;
+    VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
+      using input_t = scalar_t;
+      VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
+        using out_t = scalar_t;
+        mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
+            state.data_ptr<state_t>(), stride_state_n, stride_state_h,
+            stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
+            dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
+            B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
+            stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
+            z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
+            dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
+            out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
+            dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
+            nheads, ngroups, dim, dstate, dt_softplus);
+      });
+    });
+  });
+}
+
+// ---------------------------------------------------------------------------
+// mamba_chunk_scan_fwd_cpu
+// ---------------------------------------------------------------------------
+void mamba_chunk_scan_fwd_cpu_impl(
+    at::Tensor& out,  // [seqlen, nheads, headdim] — pre-allocated by caller
+    at::Tensor&
+        final_states,     // [batch, nheads, headdim, dstate] float32 contiguous
+    const at::Tensor& x,  // [seqlen, nheads, headdim]
+    const at::Tensor&
+        dt,  // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
+    const at::Tensor& A,                 // [nheads] float32
+    const at::Tensor& B,                 // [seqlen, ngroups, dstate]
+    const at::Tensor& C,                 // [seqlen, ngroups, dstate]
+    const c10::optional<at::Tensor>& D,  // [nheads] float32 (optional)
+    const c10::optional<at::Tensor>& z,  // [seqlen, nheads, headdim] (optional)
+    const at::Tensor& cu_seqlens         // [batch+1] int32
+) {
+  const at::ScalarType input_type = x.scalar_type();
+
+  auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+  at::Tensor x_in = ensure_contig(x);
+  at::Tensor B_in = ensure_contig(B);
+  at::Tensor C_in = ensure_contig(C);
+  at::Tensor z_in;
+  if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
+
+  // A and D are float32 model parameters, potentially broadcast-expanded.
+  // Strip trailing broadcast dims to get a contiguous (nheads,) array.
+  auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
+    at::Tensor r = t;
+    while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
+    if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
+    return r.is_contiguous() ? r : r.contiguous();
+  };
+  at::Tensor A_f32 = to_per_head_f32(A);
+  at::Tensor D_f32;
+  if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
+
+  // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
+  // Python.
+  at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
+  if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
+
+  at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
+
+  const int64_t batch = final_states.size(0);
+  const int64_t nheads = final_states.size(1);
+  const int64_t headdim = final_states.size(2);
+  const int64_t dstate = final_states.size(3);
+  const int64_t ngroups = B_in.size(1);
+
+  TORCH_CHECK(final_states.is_contiguous(),
+              "mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
+  TORCH_CHECK(out.is_contiguous(),
+              "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
+              "raw data_ptr)");
+
+  VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
+    mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
+        final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
+        dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
+        B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
+        D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
+        z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
+        out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
+        ngroups, headdim, dstate);
+  });
+}
@@ -0,0 +1,382 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+//
+// Fused CPU vector kernels for Mamba decode-step hotspots:
+//   - causal_conv1d_update  (depthwise 1-D conv state roll + compute)
+//   - selective_state_update (SSM recurrence, single-step)
+
+#pragma once
+
+#include "cpu_types.hpp"
+#include <cmath>
+#include <cstring>
+#include <cstdint>
+#include <algorithm>
+
+namespace mamba_cpu {
+
+// ---------------------------------------------------------------------------
+// causal_conv1d_update — templated for native BF16/FP32
+//
+// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
+// Explicit strides are passed so the kernel writes directly into the
+// correct memory locations without making a contiguous copy of the full
+// paged tensor (which was the source of the 34-41% direct_copy_kernel).
+//
+//   stride_s_slot  = state.stride(0)  — between cache slots
+//   stride_s_dim   = state.stride(1)  — between conv_dim channels
+//   stride_s_state = state.stride(2)  — between state elements
+//
+// When stride_s_state == 1 (contiguous), the memmove fast path is used.
+// ---------------------------------------------------------------------------
+template <typename scalar_t>
+inline void causal_conv1d_update_kernel(
+    const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
+    int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
+    const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
+    scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
+    int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
+    int64_t width, int64_t state_len, bool do_silu) {
+#pragma omp parallel for
+  for (int64_t b = 0; b < batch; ++b) {
+    int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
+    if (cache_idx == pad_slot_id) continue;
+
+    for (int64_t t = 0; t < seqlen; ++t) {
+      const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
+      scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
+      // Base of this slot in the (possibly non-contiguous) paged state
+      scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
+
+      for (int64_t d = 0; d < dim; ++d) {
+        float x_val = static_cast<float>(x_b[d * seqlen]);
+        scalar_t* sd = s_base + d * stride_s_dim;  // start of this dim's state
+        const scalar_t* w = weight_ptr + d * width;
+
+        // Accumulate in float32 for precision
+        float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
+        for (int64_t k = 0; k < state_len; ++k) {
+          acc += static_cast<float>(w[k]) *
+                 static_cast<float>(sd[k * stride_s_state]);
+        }
+        acc += static_cast<float>(w[state_len]) * x_val;
+
+        // Shift state left and append new input.
+        // Use memmove when contiguous (stride==1); element loop otherwise.
+        if (stride_s_state == 1) {
+          if (state_len > 1)
+            std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
+          if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
+        } else {
+          for (int64_t k = 0; k < state_len - 1; ++k)
+            sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
+          if (state_len > 0)
+            sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
+        }
+
+        if (do_silu) {
+          float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
+                                     : std::exp(acc) / (1.0f + std::exp(acc));
+          acc *= sigmoid;
+        }
+        out_b[d * seqlen] = static_cast<scalar_t>(acc);
+      }
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// selective_state_update
+//
+// Template parameters:
+//   state_t  - dtype of ssm_state cache (typically BFloat16)
+//   input_t  - dtype of x, B, C (typically BFloat16)
+//   out_t    - dtype of output tensor (typically BFloat16)
+//              Write directly — no float32 intermediate buffer needed.
+//
+// A, D, dt_bias are accepted as const float* (they are always float32
+// model parameters in Mamba2).  This eliminates the per-call float32→BF16
+// conversion and the .contiguous() materialisation of the broadcast-expand.
+//
+// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
+// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
+// ---------------------------------------------------------------------------
+template <typename state_t, typename input_t, typename out_t = float>
+inline void selective_state_update_kernel(
+    state_t* __restrict__ state_ptr, int64_t stride_state_n,
+    int64_t stride_state_h, int64_t stride_state_d,
+    const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
+    // dt: (N, nheads) — scalar per head, NOT expanded to head_dim
+    const float* __restrict__ dt_ptr, int64_t stride_dt_n,
+    // A: (nheads,) float32 — scalar per head
+    const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
+    const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
+    // D: (nheads,) float32 — scalar per head (nullptr if not used)
+    const float* __restrict__ D_ptr,
+    // z: same shape as x (optional)
+    const input_t* __restrict__ z_ptr,
+    // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
+    const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
+    int64_t stride_out_n, int64_t stride_out_h,
+    const int32_t* __restrict__ state_batch_indices,
+    const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
+    const int32_t* __restrict__ num_accepted_tokens,
+    const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
+    int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
+  using state_vec_t = vec_op::vec_t<state_t>;
+  using input_vec_t = vec_op::vec_t<input_t>;
+  constexpr int VEC_ELEM_NUM = 8;
+
+  int64_t nheads_per_group = nheads / ngroups;
+
+  for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
+    int64_t bos, seq_len;
+    if (cu_seqlens != nullptr) {
+      bos = cu_seqlens[seq_idx];
+      seq_len = cu_seqlens[seq_idx + 1] - bos;
+    } else {
+      bos = seq_idx;
+      seq_len = 1;
+    }
+
+    int64_t state_read_idx = (state_batch_indices != nullptr)
+                                 ? state_batch_indices[seq_idx]
+                                 : seq_idx;
+    if (state_read_idx == null_block_id) continue;
+
+    int64_t state_write_idx = (num_accepted_tokens == nullptr)
+                                  ? ((dst_state_batch_indices != nullptr)
+                                         ? dst_state_batch_indices[seq_idx]
+                                         : state_read_idx)
+                                  : -1;
+
+    state_t* s = state_ptr + state_read_idx * stride_state_n;
+
+    for (int64_t t = 0; t < seq_len; ++t) {
+      int64_t token_idx = bos + t;
+      const input_t* x_tok = x_ptr + token_idx * stride_x_n;
+      // dt: (N, nheads) — one float per head per token
+      const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
+      const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
+      const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
+      out_t* out_tok = out_ptr + token_idx * stride_out_n;
+
+#pragma omp parallel for
+      for (int64_t h = 0; h < nheads; ++h) {
+        int64_t g = h / nheads_per_group;
+        const input_t* x_h = x_tok + h * stride_x_h;
+        const input_t* B_g = B_tok + g * stride_BC_g;
+        const input_t* C_g = C_tok + g * stride_BC_g;
+        out_t* out_h = out_tok + h * stride_out_h;
+        state_t* s_h = s + h * stride_state_h;
+
+        // Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
+        float dt_val = dt_tok[h];
+        if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
+        if (dt_softplus) {
+          dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
+        }
+        const float A_val = A_ptr[h];  // scalar: same for all dim, dstate
+        const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
+
+        const input_t* z_h =
+            (z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
+                               : nullptr;
+
+        vec_op::FP32Vec8 dt_vec(dt_val);
+        // dA = exp(A * dt): A and dt are SCALARS per head, so compute once
+        // and broadcast.  This saves 7 redundant std::exp() calls that
+        // FP32Vec8::exp() would otherwise make on the broadcast vector.
+        const float dA_scalar = std::exp(A_val * dt_val);
+        vec_op::FP32Vec8 dA(dA_scalar);  // broadcast
+
+        for (int64_t d = 0; d < dim; ++d) {
+          float x_val = static_cast<float>(x_h[d]);
+
+          vec_op::FP32Vec8 out_vec(0.0f);
+          state_t* s_hd = s_h + d * stride_state_d;
+          const input_t* B_g_base = B_g;
+          const input_t* C_g_base = C_g;
+
+          vec_op::FP32Vec8 x_vec(x_val);
+          // dBx = B * x * dt — same dA for all dstate (A is scalar)
+          // s_new = s * dA + B * x * dt
+
+          int64_t n = 0;
+          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
+            vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
+            vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
+            vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
+
+            vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
+            vec_op::FP32Vec8 s_new = s_v * dA + dBx;
+
+            state_vec_t(s_new).save(s_hd + n);
+            out_vec = out_vec + s_new * C_v;
+          }
+
+          float out_val = out_vec.reduce_sum();
+          for (; n < dstate; ++n) {
+            // Reuse dA_scalar computed once per head — no exp() re-call
+            float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
+            float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
+            s_hd[n] = static_cast<state_t>(s_new);
+            out_val += s_new * static_cast<float>(C_g[n]);
+          }
+
+          if (D_ptr != nullptr) out_val += x_val * D_val;
+          if (z_h != nullptr) {
+            float z_val = static_cast<float>(z_h[d]);
+            float sigmoid = (z_val >= 0)
+                                ? 1.0f / (1.0f + std::exp(-z_val))
+                                : std::exp(z_val) / (1.0f + std::exp(z_val));
+            out_val *= z_val * sigmoid;
+          }
+          out_h[d] = static_cast<out_t>(out_val);
+        }
+      }
+
+      if (num_accepted_tokens != nullptr &&
+          dst_state_batch_indices != nullptr) {
+        int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
+        if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
+          state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
+          std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
+        }
+      }
+    }
+
+    if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
+        state_write_idx != state_read_idx) {
+      state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
+      std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// mamba_chunk_scan_fwd
+//
+// Prefill SSM recurrence for Mamba2 / SSD models.
+//
+// Key difference from selective_state_update_kernel (decode path):
+//   - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
+//     Each thread owns a (batch, head) slice and runs the entire token
+//     sequence without any per-token OpenMP synchronisation overhead.
+//     For seqlen=256, this eliminates 256 thread-barrier launches per batch.
+//
+// `dt` arrives already processed (float32, after bias + softplus + clamp)
+// to keep this kernel simple. Preprocessing is done in the Python wrapper.
+//
+// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
+// tensor, pre-initialised by the caller (zero or from initial_states).
+// Each (b, h) slice is private to exactly one thread via collapse(2), so
+// there are no write conflicts.
+//
+// D is treated as a scalar per head ([nheads] float32).
+// ---------------------------------------------------------------------------
+template <typename input_t>
+inline void mamba_chunk_scan_fwd_kernel(
+    float* __restrict__ states_ptr,     // [batch, nheads, headdim, dstate] f32
+    const input_t* __restrict__ x_ptr,  // [seqlen, nheads, headdim]
+    const float* __restrict__ dt_ptr,   // [seqlen, nheads] f32 (preprocessed)
+    const float* __restrict__ A_ptr,    // [nheads] f32
+    const input_t* __restrict__ B_ptr,  // [seqlen, ngroups, dstate]
+    const input_t* __restrict__ C_ptr,  // [seqlen, ngroups, dstate]
+    const float* __restrict__ D_ptr,    // [nheads] f32 (nullable)
+    const input_t* __restrict__ z_ptr,  // [seqlen, nheads, headdim] (nullable)
+    input_t* __restrict__ out_ptr,      // [seqlen, nheads, headdim]
+    const int32_t* __restrict__ cu_seqlens,  // [batch+1] int32
+    int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
+    int64_t dstate) {
+  using input_vec_t = vec_op::vec_t<input_t>;
+  constexpr int VEC_ELEM_NUM = 8;
+
+  const int64_t nheads_per_group = nheads / ngroups;
+  // states layout: [batch, nheads, headdim, dstate] contiguous (caller
+  // guarantee)
+  const int64_t stride_s_b = nheads * headdim * dstate;
+  const int64_t stride_s_h = headdim * dstate;
+  // stride_s_d = dstate, stride_s_n = 1
+
+#pragma omp parallel for collapse(2) schedule(static)
+  for (int64_t b = 0; b < batch; ++b) {
+    for (int64_t h = 0; h < nheads; ++h) {
+      const int64_t seq_start = cu_seqlens[b];
+      const int64_t seq_end = cu_seqlens[b + 1];
+      const int64_t g = h / nheads_per_group;
+
+      const float A_val = A_ptr[h];
+      const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
+
+      // Working state slice: states[b, h, :, :] — float32, headdim * dstate.
+      // Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
+      float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
+
+      for (int64_t t = seq_start; t < seq_end; ++t) {
+        const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
+        const float* dt_h = dt_ptr + t * nheads + h;
+        const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
+        const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
+        const input_t* z_h = (z_ptr != nullptr)
+                                 ? z_ptr + t * nheads * headdim + h * headdim
+                                 : nullptr;
+        input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
+
+        const float dt_val = *dt_h;
+        const float dA_val = std::exp(A_val * dt_val);
+        const vec_op::FP32Vec8 dA_vec(dA_val);  // broadcast scalar
+        const vec_op::FP32Vec8 dt_vec(dt_val);
+
+        for (int64_t d = 0; d < headdim; ++d) {
+          const float x_val = static_cast<float>(x_h[d]);
+          float* s_bhd = s_bh + d * dstate;  // [dstate] contiguous float32
+
+          // Vectorised SSM update + readout over dstate:
+          //   s_new = s * dA + x * dt * B
+          //   y    += s_new * C
+          int64_t n = 0;
+          vec_op::FP32Vec8 y_vec(0.0f);
+          const vec_op::FP32Vec8 x_vec(x_val);
+
+          for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
+            const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
+            const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
+            const vec_op::FP32Vec8 s_v(s_bhd + n);
+
+            const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
+            s_new.save(s_bhd + n);
+            y_vec = y_vec + s_new * C_v;
+          }
+
+          float y_val = y_vec.reduce_sum();
+
+          // Scalar tail for remaining dstate elements
+          for (; n < dstate; ++n) {
+            const float B_n = static_cast<float>(B_g[n]);
+            const float C_n = static_cast<float>(C_g[n]);
+            const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
+            s_bhd[n] = s_new;
+            y_val += s_new * C_n;
+          }
+
+          // D skip connection (scalar per head)
+          if (D_ptr != nullptr) y_val += x_val * D_val;
+
+          // z gating: out = y * z * sigmoid(z)  (SiLU)
+          if (z_h != nullptr) {
+            const float z_val = static_cast<float>(z_h[d]);
+            const float sigmoid =
+                (z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
+                                : std::exp(z_val) / (1.0f + std::exp(z_val));
+            y_val *= z_val * sigmoid;
+          }
+
+          out_h[d] = static_cast<input_t>(y_val);
+        }
+      }
+    }
+  }
+}
+
+}  // namespace mamba_cpu
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
                                      torch::Tensor slot_mapping,
                                      const int64_t block_size);

+at::Tensor causal_conv1d_update_cpu_impl(
+    at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
+    const c10::optional<at::Tensor>& bias,
+    const c10::optional<std::string>& activation,
+    const c10::optional<at::Tensor>& conv_state_indices,
+    const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
+
+void selective_state_update_cpu_impl(
+    at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
+    const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
+    const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
+    const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
+    const c10::optional<at::Tensor>& state_batch_indices,
+    const c10::optional<at::Tensor>& dst_state_batch_indices,
+    int64_t null_block_id, at::Tensor& out,
+    const c10::optional<at::Tensor>& num_accepted_tokens,
+    const c10::optional<at::Tensor>& cu_seqlens);
+
+void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
+                                   const at::Tensor& x, const at::Tensor& dt,
+                                   const at::Tensor& A, const at::Tensor& B,
+                                   const at::Tensor& C,
+                                   const c10::optional<at::Tensor>& D,
+                                   const c10::optional<at::Tensor>& z,
+                                   const at::Tensor& cu_seqlens);
+
 void init_cpu_memory_env(std::vector<int64_t> node_ids);

 namespace cpu_utils {
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "block_size) -> ()",
      &compute_slot_mapping_kernel_impl);

+  // Mamba CPU kernels
+  ops.def(
+      "causal_conv1d_update_cpu_vec("
+      "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
+      "Tensor? bias, str? activation, Tensor? conv_state_indices, "
+      "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
+      &causal_conv1d_update_cpu_impl);
+
+  ops.def(
+      "selective_state_update_cpu("
+      "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
+      "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
+      "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
+      "SymInt null_block_id, Tensor(a13!) out, "
+      "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
+      &selective_state_update_cpu_impl);
+
+  ops.def(
+      "mamba_chunk_scan_fwd_cpu("
+      "Tensor(a0!) out, Tensor(a1!) final_states, "
+      "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
+      "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
+      &mamba_chunk_scan_fwd_cpu_impl);
+
  ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);

  // Speculative decoding kernels
@@ -47,6 +47,8 @@ torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,
                                   torch::stable::Tensor const& perm);

 #ifndef USE_ROCM
+std::string get_compiled_cuda_archs();
+
 bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
 bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
 bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
    // AWQ zeros: (size_k // group_size, size_n // 8)
    const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
    int32_t group_size) {
-  int32_t val =
-      qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
-  int32_t zero =
-      qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
-             blockIdx.y];
+  // Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
+  // blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
+  // same row (stride 1) instead of striding across rows (stride size_n/8).
+  int col = blockIdx.y * 32 + threadIdx.x;
+  if (col >= size_n / 8) return;
+  (void)size_k;
+
+  int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
+  int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
  int32_t new_val = 0;

 #pragma unroll
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
    zero >>= 4;
  }

-  output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
+  output[blockIdx.x * (size_n / 8) + col] = new_val;
 }

 torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
                    "qweight.size(0) % qzeros.size(0) != 0");
    STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");

-    dim3 blocks(size_k / 32, size_n / 8);
+    dim3 blocks(size_k, (size_n / 8 + 31) / 32);
    marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
        reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
        reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
@@ -51,7 +51,8 @@ void cutlass_moe_mm_sm90(torch::stable::Tensor& out_tensors,

 #endif

-#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
+#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
+    ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
 void cutlass_moe_mm_sm100(torch::stable::Tensor& out_tensors,
                          torch::stable::Tensor const& a_tensors,
                          torch::stable::Tensor const& b_tensors,
@@ -83,8 +84,9 @@ void cutlass_scaled_mm_sm100(torch::stable::Tensor& c,
                             std::optional<torch::stable::Tensor> const& bias);
 #endif

-#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined(ENABLE_CUTLASS_MOE_SM100) && ENABLE_CUTLASS_MOE_SM100) || \
+#if (defined(ENABLE_CUTLASS_MOE_SM90) && ENABLE_CUTLASS_MOE_SM90) || \
+    (defined(ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) &&                   \
+     ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) ||                           \
    (defined(ENABLE_CUTLASS_MOE_SM120) && ENABLE_CUTLASS_MOE_SM120)
 void get_cutlass_moe_mm_data_caller(
    const torch::stable::Tensor& topk_ids,
@@ -175,11 +177,14 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {

 bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
  // CUTLASS grouped FP8 kernels need at least CUDA 12.3 and SM90 (Hopper)
-  // or CUDA 12.8 and SM100 (Blackwell). Only report archs that have an
-  // actual cutlass_moe_mm dispatch compiled into this file.
+  // or CUDA 12.8 and SM10x/SM11x (Blackwell / Thor). CUDA 12 reports Thor as
+  // SM101 while CUDA 13 reports it as SM110, but both use this sm100-named
+  // implementation. Only report archs that have an actual cutlass_moe_mm
+  // dispatch compiled into this file.

 #if defined CUDA_VERSION
-  #if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
+  #if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
+      ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
  if (cuda_device_capability >= 100 && cuda_device_capability < 120) {
    return CUDA_VERSION >= 12080;
  }
@@ -281,8 +286,11 @@ void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
                    torch::stable::Tensor const& c_strides, bool per_act_token,
                    bool per_out_ch) {
  int32_t version_num = get_sm_version_num();
-#if defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100
-  if (version_num >= 100 && version_num < 110) {
+#if defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X && \
+    ENABLE_CUTLASS_MOE_SM10X_OR_SM11X
+  // Keep runtime dispatch aligned with the CMake arch list and support query:
+  // CUDA 12 Thor is SM101 and CUDA 13 Thor is SM110.
+  if (version_num >= 100 && version_num < 120) {
    cutlass_moe_mm_sm100(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
                         expert_offsets, problem_sizes, a_strides, b_strides,
                         c_strides, per_act_token, per_out_ch);
@@ -316,8 +324,9 @@ void get_cutlass_moe_mm_data(
  // This function currently gets compiled only if we have a valid cutlass moe
  // mm to run it for.
  int32_t version_num = get_sm_version_num();
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
+    (defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X &&                   \
+     ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) ||                          \
    (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
  get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
                                 problem_sizes2, input_permutation,
@@ -338,8 +347,9 @@ void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
    torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
    const bool swap_ab) {
  int32_t version_num = get_sm_version_num();
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
+    (defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X &&                   \
+     ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) ||                          \
    (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
  get_cutlass_moe_mm_problem_sizes_from_expert_offsets_caller(
      expert_first_token_offset, problem_sizes1, problem_sizes2, n, k, swap_ab);
@@ -362,8 +372,9 @@ void get_cutlass_batched_moe_mm_data(
  // This function currently gets compiled only if we have a valid cutlass moe
  // mm to run it for.
  int32_t version_num = get_sm_version_num();
-#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) ||   \
-    (defined ENABLE_CUTLASS_MOE_SM100 && ENABLE_CUTLASS_MOE_SM100) || \
+#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
+    (defined ENABLE_CUTLASS_MOE_SM10X_OR_SM11X &&                   \
+     ENABLE_CUTLASS_MOE_SM10X_OR_SM11X) ||                          \
    (defined ENABLE_CUTLASS_MOE_SM120 && ENABLE_CUTLASS_MOE_SM120)
  get_cutlass_batched_moe_mm_data_caller(expert_offsets, problem_sizes1,
                                         problem_sizes2, expert_num_tokens,
@@ -4,6 +4,10 @@

 #include <torch/csrc/stable/library.h>

+#ifndef USE_ROCM
+std::string get_compiled_cuda_archs() { return VLLM_COMPILED_CUDA_ARCHS; }
+#endif
+
 // Register ops with STABLE_TORCH_LIBRARY for libtorch stable ABI compatibility.
 // Note: We register under namespace "_C" so ops are accessible as
 // torch.ops._C.<op_name> for compatibility with existing code.
@@ -30,6 +34,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
  ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");

  ops.def("get_cuda_view_from_cpu_tensor(Tensor cpu_tensor) -> Tensor");
+  ops.def("get_compiled_cuda_archs() -> str");

 #ifndef USE_ROCM

@@ -763,6 +768,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C_cuda_utils, CompositeExplicitAutograd,
 // ops.impl("op_name", &func) without a dispatch key in the non-stable API.
 STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, ops) {
 #ifndef USE_ROCM
+  ops.impl("get_compiled_cuda_archs", TORCH_BOX(&get_compiled_cuda_archs));
  ops.impl("cutlass_scaled_mm_supports_fp8",
           TORCH_BOX(&cutlass_scaled_mm_supports_fp8));
  ops.impl("cutlass_group_gemm_supported",
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
 | `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
 | `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
 | `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
-| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
+| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
 | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
 | `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
 | `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:

 - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
 - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
+- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)

 ### SiLU+Mul + Quantization (`fuse_act_quant`)

@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
 - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
 - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
 - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
+- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)

 ### RMSNorm + Padding (`fuse_act_padding`)

@@ -68,13 +68,14 @@ vllm serve <model> \
 | --- | --- | --- | --- | --- |
 | `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
 | `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
-| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
+| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
+| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
 | `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
 | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
 | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
 | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
 | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
-| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
+| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
 | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |

 ## Secondary Tiers
@@ -179,7 +180,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on

 - `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
 - For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
+- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
 - FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
 - Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.

@@ -46,6 +46,26 @@ export CPU_ARCH=$(uname -m) # x86_64 or aarch64
 uv pip install https://github.com/vllm-project/vllm/releases/download/v${VLLM_VERSION}/vllm-${VLLM_VERSION}+cu${CUDA_VERSION}-cp38-abi3-manylinux_2_28_${CPU_ARCH}.whl --extra-index-url https://download.pytorch.org/whl/cu${CUDA_VERSION}
 ```

+!!! warning "CUDA architecture coverage"
+
+    Pre-built CUDA wheels are compiled for the CUDA architectures selected by
+    vLLM's release and build pipelines. This list is intentionally smaller than
+    every architecture CMake can build, because each additional architecture
+    increases wheel size.
+
+    In particular, CUDA 12.9 wheels do not use CUDA 13 family-specific targets.
+    To keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
+    architecture-specific Blackwell/Thor targets, such as `sm_103` or
+    `sm_121`, even though vLLM can build them from source. CUDA 13 wheels use
+    family-specific targets such as `sm_100f`, `sm_110f`, and `sm_120f`, which
+    cover the corresponding major-version GPU family.
+
+    If vLLM logs a warning that your visible CUDA device is not covered by the
+    wheel's compiled CUDA architectures, or if you see a CUDA error such as
+    `no kernel image is available for execution on the device`, install a CUDA
+    13 wheel when possible, or build from source with a `TORCH_CUDA_ARCH_LIST`
+    that includes your GPU.
+
 #### Install the latest code

 LLM inference is a fast-evolving field, and the latest code may contain bug fixes, performance improvements, and new features that are not released yet. To allow users to try the latest code without waiting for the next release, vLLM provides wheels for every commit since `v0.5.3` on <https://wheels.vllm.ai/nightly>. There are multiple indices that could be used:
@@ -31,10 +31,8 @@
 | THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
 | chuhac/TeleChat2-35B                               | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅                    |                        |       |
 | 01-ai/Yi1.5-34B-Chat                               | YiForCausalLM                                    | ✅                    |                        |       |
-| THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
 | deepseek-ai/DeepSeek-Coder-33B-base                | DeepSeekCoderForCausalLM                         | ✅                    |                        |       |
 | meta-llama/Llama-2-13b-chat-hf                     | LlamaForCausalLM                                 | ✅                    |                        |       |
-| THUDM/CodeGeex4-All-9B                             | CodeGeexForCausalLM                              | ✅                    |                        |       |
 | Qwen/Qwen1.5-14B-Chat                              | QwenForCausalLM                                  | ✅                    |                        |       |
 | Qwen/Qwen1.5-32B-Chat                              | QwenForCausalLM                                  | ✅                    |                        |       |
 | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic    | LlamaForCausalLM                                 |                       | ✅                     |       |
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
 - `response_format`: Format of the response ("json", "text") (optional)
 - `temperature`: Sampling temperature between 0 and 1 (optional)

-For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
+For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).

 **Response Format:**

@@ -174,6 +174,33 @@ If the script runs successfully, you should see the message `sanity check is suc

 If the test script hangs or crashes, usually it means the hardware/drivers are broken in some sense. You should try to contact your system administrator or hardware vendor for further assistance. As a common workaround, you can try to tune some NCCL environment variables, such as `export NCCL_P2P_DISABLE=1` to see if it helps. Please check [their documentation](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html) for more information. Please only use these environment variables as a temporary workaround, as they might affect the performance of the system. The best solution is still to fix the hardware/drivers so that the test script can run successfully.

+## CUDA architecture not covered by the wheel
+
+If vLLM logs a warning that the current wheel was built for a set of CUDA
+architectures but one of your visible CUDA devices is not covered, the installed
+wheel may not contain native CUDA kernels for that GPU. The same issue may also
+surface later as a CUDA runtime error such as `no kernel image is available for
+execution on the device`.
+
+The architecture list in a pre-built wheel is determined by the vLLM release and
+build pipelines, then filtered by CMake before individual kernels choose their
+own per-kernel architectures. It is not the same as the full set of
+architectures that vLLM can build from source.
+
+CUDA 12.9 wheels use architecture-specific targets for newer NVIDIA GPUs. To
+keep wheel size bounded, published CUDA 12.9 wheels may omit some newer
+Blackwell/Thor targets such as `sm_103` or `sm_121`, even though vLLM can build
+them from source. CUDA 13 wheels can use family-specific targets such as
+`sm_100f`, `sm_110f`, and `sm_120f`, which cover the corresponding
+major-version GPU family.
+
+If you see this warning or error, install a CUDA 13 wheel when possible.
+Otherwise, build vLLM from source and set `TORCH_CUDA_ARCH_LIST` to include
+your GPU's compute capability. See the CUDA installation guide's
+[pre-built wheels](../getting_started/installation/gpu.md#pre-built-wheels) and
+[build wheel from source](../getting_started/installation/gpu.md#build-wheel-from-source)
+sections for details.
+
 ## Python multiprocessing

 ### `RuntimeError` Exception
@@ -5560,6 +5560,7 @@ dependencies = [
 "tracing",
 "tracing-subscriber",
 "uuid",
+ "vllm-bench",
 "vllm-chat",
 "vllm-engine-core-client",
 "vllm-managed-engine",
@@ -6320,9 +6321,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"

 [[package]]
 name = "xgrammar-structural-tag"
-version = "0.1.0+xgrammar.0.2.2.4d145cc"
+version = "0.2.0+xgrammar.0.2.4.dd729e7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
+checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
 dependencies = [
 "auto_impl",
 "serde",
@@ -132,6 +132,7 @@ trait-set = "0.3.0"
 url = "2.5.7"
 uuid = { version = "1.22.0", features = ["v4"] }
 validator = { version = "0.20.0", features = ["derive"] }
+vllm-bench = { path = "src/bench" }
 vllm-chat = { path = "src/chat" }
 vllm-engine-core-client = { path = "src/engine-core-client" }
 vllm-llm = { path = "src/llm" }
@@ -142,7 +143,7 @@ vllm-server = { path = "src/server" }
 vllm-text = { path = "src/text" }
 vllm-tokenizer = { path = "src/tokenizer" }
 winnow = { version = "1.0.2", features = ["simd"] }
-xgrammar-structural-tag = "0.1.0"
+xgrammar-structural-tag = "0.2.0"
 zeromq = { version = "0.6.0", default-features = false, features = [
    "tokio-runtime",
    "all-transport",
@@ -3,8 +3,6 @@

 use std::fmt;

-use clap::Parser;
-
 /// Backend type for the benchmark endpoint.
 #[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
 pub enum BackendKind {
@@ -77,7 +75,7 @@ pub enum DatasetName {
    ShareGpt,
    #[value(name = "sonnet")]
    Sonnet,
-    #[value(name = "speed-bench")]
+    #[value(name = "speed-bench", alias = "speed_bench")]
    SpeedBench,
    #[value(name = "hf")]
    Hf,
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
 }

 /// High-performance benchmark client for vLLM serving endpoints.
-#[derive(Parser, Debug, Clone)]
-#[command(
-    name = "vllm-bench",
-    about = "Benchmark online serving throughput",
-    version
-)]
-pub struct Cli {
+#[derive(clap::Args, Debug, Clone)]
+pub struct BenchServeArgs {
    /// The type of backend or endpoint to use for the benchmark.
    #[arg(long, default_value = "openai")]
    pub backend: BackendKind,
@@ -659,7 +652,7 @@ pub struct Cli {
    pub lora_assignment: LoraAssignment,
 }

-impl Cli {
+impl BenchServeArgs {
    /// Resolve the base URL from explicit --base-url or from --host/--port.
    pub fn resolve_base_url(&self) -> String {
        if let Some(ref base) = self.base_url {
@@ -4,7 +4,9 @@
 use std::collections::HashMap;
 use std::sync::Arc;

-use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
+use crate::cli::{
+    BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
+};
 use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
 use crate::error::{BenchError, Result};

@@ -215,63 +217,63 @@ pub struct BenchConfig {
 }

 impl BenchConfig {
-    pub fn from_cli(cli: &Cli) -> Result<Self> {
-        if cli.burstiness <= 0.0 {
+    pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
+        if args.burstiness <= 0.0 {
            return Err(BenchError::Config("Burstiness must be positive".into()));
        }

-        if cli.num_prompts == 0 {
+        if args.num_prompts == 0 {
            return Err(BenchError::Config(
                "--num-prompts must be at least 1".into(),
            ));
        }

-        if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
+        if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
            return Err(BenchError::Config(
                "--request-rate must be positive (or inf)".into(),
            ));
        }
-        if cli.max_model_len == Some(0) {
+        if args.max_model_len == Some(0) {
            return Err(BenchError::Config(
                "--max-model-len must be at least 1".into(),
            ));
        }

-        let base_url = cli.resolve_base_url();
-        let api_url = cli.resolve_api_url();
+        let base_url = args.resolve_base_url();
+        let api_url = args.resolve_api_url();

-        let extra_headers = cli.parse_headers()?;
-        let mut extra_body = cli.parse_extra_body()?;
+        let extra_headers = args.parse_headers()?;
+        let mut extra_body = args.parse_extra_body()?;

        // Merge sampling parameters into extra_body (matches Python behavior).
        // Python collects non-None sampling params and merges them UNDER extra_body,
        // meaning extra_body keys take precedence over sampling params.
        {
            let mut sampling_params = serde_json::Map::new();
-            if let Some(v) = cli.top_p {
+            if let Some(v) = args.top_p {
                sampling_params.insert("top_p".into(), serde_json::json!(v));
            }
-            if let Some(v) = cli.top_k {
+            if let Some(v) = args.top_k {
                sampling_params.insert("top_k".into(), serde_json::json!(v));
            }
-            if let Some(v) = cli.min_p {
+            if let Some(v) = args.min_p {
                sampling_params.insert("min_p".into(), serde_json::json!(v));
            }
-            if let Some(v) = cli.temperature {
+            if let Some(v) = args.temperature {
                sampling_params.insert("temperature".into(), serde_json::json!(v));
            }
-            if let Some(v) = cli.frequency_penalty {
+            if let Some(v) = args.frequency_penalty {
                sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
            }
-            if let Some(v) = cli.presence_penalty {
+            if let Some(v) = args.presence_penalty {
                sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
            }
-            if let Some(v) = cli.repetition_penalty {
+            if let Some(v) = args.repetition_penalty {
                sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
            }

            if !sampling_params.is_empty() {
-                if !cli.backend.is_openai_compatible() {
+                if !args.backend.is_openai_compatible() {
                    return Err(BenchError::Config(
                        "Sampling parameters are only supported by openai-compatible backends."
                            .into(),
@@ -299,7 +301,7 @@ impl BenchConfig {
        }

        // Parse metadata
-        let metadata = match &cli.metadata {
+        let metadata = match &args.metadata {
            None => None,
            Some(items) => {
                let mut pairs = Vec::new();
@@ -314,24 +316,24 @@ impl BenchConfig {
        };

        // Parse goodput SLOs
-        let goodput = parse_goodput(&cli.goodput)?;
+        let goodput = parse_goodput(&args.goodput)?;

        // Parse ramp-up config
-        let ramp_up = parse_ramp_up(cli)?;
+        let ramp_up = parse_ramp_up(args)?;

        // Default percentile metrics based on backend type
-        let default_percentile_metrics = if cli.backend.is_pooling() {
+        let default_percentile_metrics = if args.backend.is_pooling() {
            "e2el"
        } else {
            "ttft,tpot,itl,e2el"
        };
        let percentile_metrics_str =
-            cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
+            args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
        let selected_percentile_metrics: Vec<String> =
            percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();

-        let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
-        let sweep_summary_percentiles = cli
+        let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
+        let sweep_summary_percentiles = args
            .sweep_summary_percentiles
            .as_deref()
            .map(|raw| parse_percentiles(raw, true))
@@ -344,38 +346,38 @@ impl BenchConfig {
            selected_percentiles.push(90.0);
        }

-        let tokenizer_id = if cli.skip_tokenizer_init {
+        let tokenizer_id = if args.skip_tokenizer_init {
            None
        } else {
-            Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
+            args.tokenizer.clone().or_else(|| args.model.clone())
        };

        // Resolve input/output lengths
-        let random_input_len = cli.resolved_random_input_len();
-        let random_output_len = cli.resolved_random_output_len();
-        let per_turn_input_len = cli.resolved_per_turn_input_len();
+        let random_input_len = args.resolved_random_input_len();
+        let random_output_len = args.resolved_random_output_len();
+        let per_turn_input_len = args.resolved_per_turn_input_len();

        // Normalized multi-turn turn counts (computed in validation block below, defaults
        // to num_turns if multi-turn mode is not active)
-        let mut multi_turn_min_turns = cli.multi_turn_num_turns;
-        let mut multi_turn_max_turns = cli.multi_turn_num_turns;
+        let mut multi_turn_min_turns = args.multi_turn_num_turns;
+        let mut multi_turn_max_turns = args.multi_turn_num_turns;

        // For random datasets with openai-compatible backends, default to ignore_eos.
        // Exception: multi-turn mode, where ignore_eos causes unbounded context growth
        // across turns. Multi-turn uses min_tokens instead for output length control.
        // Pooling backends don't generate tokens, so ignore_eos is irrelevant.
-        let ignore_eos = if cli.backend.is_pooling() {
+        let ignore_eos = if args.backend.is_pooling() {
            false
        } else {
-            cli.ignore_eos
-                || ((cli.dataset_name == DatasetName::Random
-                    || cli.dataset_name == DatasetName::RandomMm)
-                    && cli.backend.is_openai_compatible()
-                    && !cli.multi_turn)
+            args.ignore_eos
+                || ((args.dataset_name == DatasetName::Random
+                    || args.dataset_name == DatasetName::RandomMm)
+                    && args.backend.is_openai_compatible()
+                    && !args.multi_turn)
        };

        // Pooling backends don't support multi-turn
-        if cli.backend.is_pooling() && cli.multi_turn {
+        if args.backend.is_pooling() && args.multi_turn {
            return Err(BenchError::Config(
                "Pooling/embedding backends do not support --multi-turn".into(),
            ));
@@ -383,7 +385,7 @@ impl BenchConfig {

        // LoRA validation. Adapter names must be non-empty after trim; pooling
        // backends are out of scope (vLLM LoRA routing is for generative paths).
-        let lora_modules = match cli.lora_modules.as_ref() {
+        let lora_modules = match args.lora_modules.as_ref() {
            None => None,
            Some(names) => {
                if names.is_empty() {
@@ -391,7 +393,7 @@ impl BenchConfig {
                        "--lora-modules requires at least one adapter name".into(),
                    ));
                }
-                if cli.backend.is_pooling() {
+                if args.backend.is_pooling() {
                    return Err(BenchError::Config(
                        "--lora-modules is not supported for pooling/embedding backends".into(),
                    ));
@@ -411,18 +413,18 @@ impl BenchConfig {
        };

        // Random-MM validation and config parsing
-        let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
-            if cli.backend != BackendKind::OpenaiChat {
+        let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
+            if args.backend != BackendKind::OpenaiChat {
                return Err(BenchError::Config(
                    "Multi-modal content (images) is only supported on 'openai-chat' backend."
                        .into(),
                ));
            }
            let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
-                &cli.random_mm_limit_mm_per_prompt,
+                &args.random_mm_limit_mm_per_prompt,
            )?;
            let buckets =
-                crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
+                crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
            (limit, buckets)
        } else {
            (MmLimitPerPrompt::default(), Vec::new())
@@ -432,18 +434,18 @@ impl BenchConfig {
        // sonnet (uses built-in Shakespeare's sonnets).

        // Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
-        let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
+        let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;

        // Batched inputs only make sense for pooling backends (the generation
        // backends send one prompt per request).
-        if cli.random_batch_size == 0 {
+        if args.random_batch_size == 0 {
            return Err(BenchError::Config(
                "--random-batch-size must be at least 1".into(),
            ));
        }
-        if cli.random_batch_size > 1
-            && !cli.backend.is_pooling()
-            && cli.dataset_name != DatasetName::RandomRerank
+        if args.random_batch_size > 1
+            && !args.backend.is_pooling()
+            && args.dataset_name != DatasetName::RandomRerank
        {
            return Err(BenchError::Config(
                "--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -451,16 +453,16 @@ impl BenchConfig {
        }

        // random-rerank validation (mirrors Python RandomDatasetForReranking)
-        let is_reranker = !cli.no_reranker;
-        if cli.dataset_name == DatasetName::RandomRerank {
-            if !cli.backend.is_pooling() {
+        let is_reranker = !args.no_reranker;
+        if args.dataset_name == DatasetName::RandomRerank {
+            if !args.backend.is_pooling() {
                return Err(BenchError::Config(
                    "--dataset-name random-rerank requires an embeddings/pooling backend \
                     (e.g. --backend vllm-rerank)"
                        .into(),
                ));
            }
-            if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
+            if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
                return Err(BenchError::Config(
                    "--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
                     (the query is folded into the first batch slot)"
@@ -470,8 +472,8 @@ impl BenchConfig {
        }

        // Custom dataset validation
-        if cli.dataset_name == DatasetName::Custom {
-            match cli.dataset_path.as_deref() {
+        if args.dataset_name == DatasetName::Custom {
+            match args.dataset_path.as_deref() {
                None => {
                    return Err(BenchError::Config(
                        "--dataset-path is required for --dataset-name custom \
@@ -486,7 +488,7 @@ impl BenchConfig {
                }
                _ => {}
            }
-            if !cli.skip_chat_template {
+            if !args.skip_chat_template {
                eprintln!(
                    "NOTE: client-side chat template rendering is not supported; custom \
                     dataset prompts are sent raw (equivalent to --skip-chat-template)."
@@ -495,29 +497,29 @@ impl BenchConfig {
        }

        // Prefix repetition validation
-        if cli.dataset_name == DatasetName::PrefixRepetition {
-            if cli.prefix_repetition_num_prefixes == 0 {
+        if args.dataset_name == DatasetName::PrefixRepetition {
+            if args.prefix_repetition_num_prefixes == 0 {
                return Err(BenchError::Config(
                    "--prefix-repetition-num-prefixes must be at least 1".into(),
                ));
            }
-            if cli.num_prompts < cli.prefix_repetition_num_prefixes {
+            if args.num_prompts < args.prefix_repetition_num_prefixes {
                return Err(BenchError::Config(format!(
                    "--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
-                    cli.num_prompts, cli.prefix_repetition_num_prefixes
+                    args.num_prompts, args.prefix_repetition_num_prefixes
                )));
            }
        }

        // HF dataset validation
-        if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
+        if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
            return Err(BenchError::Config(
                "--dataset-path is required for --dataset-name hf \
                 (set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
                    .into(),
            ));
        }
-        if let Some(len) = cli.hf_output_len
+        if let Some(len) = args.hf_output_len
            && len == 0
        {
            return Err(BenchError::Config(
@@ -526,13 +528,13 @@ impl BenchConfig {
        }

        // Multi-turn validation
-        if cli.multi_turn {
-            if cli.backend != BackendKind::OpenaiChat {
+        if args.multi_turn {
+            if args.backend != BackendKind::OpenaiChat {
                return Err(BenchError::Config(
                    "--multi-turn requires --backend openai-chat".into(),
                ));
            }
-            if cli.multi_turn_num_turns == 0 {
+            if args.multi_turn_num_turns == 0 {
                return Err(BenchError::Config(
                    "--multi-turn-num-turns must be at least 1".into(),
                ));
@@ -541,18 +543,18 @@ impl BenchConfig {
            // Normalize and validate min/max turns. ShareGPT only consumes max_turns
            // (the loader walks all available turns up to the cap), so the
            // min/num/max coupling used for synthetic generation does not apply.
-            if cli.dataset_name == DatasetName::ShareGpt {
-                if cli.multi_turn_max_turns == 1 {
+            if args.dataset_name == DatasetName::ShareGpt {
+                if args.multi_turn_max_turns == 1 {
                    return Err(BenchError::Config(
                        "--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
                    ));
                }
            } else {
                (multi_turn_min_turns, multi_turn_max_turns) =
-                    match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
-                        (0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
-                        (m, 0) => (m, cli.multi_turn_num_turns),
-                        (0, x) => (cli.multi_turn_num_turns, x),
+                    match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
+                        (0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
+                        (m, 0) => (m, args.multi_turn_num_turns),
+                        (0, x) => (args.multi_turn_num_turns, x),
                        (m, x) => (m, x),
                    };
                if multi_turn_min_turns < 1 {
@@ -575,8 +577,8 @@ impl BenchConfig {
            }

            // Validate prefix sharing ratios
-            let pg = cli.multi_turn_prefix_global_ratio;
-            let pc = cli.multi_turn_prefix_conversation_ratio;
+            let pg = args.multi_turn_prefix_global_ratio;
+            let pc = args.multi_turn_prefix_conversation_ratio;
            if !(0.0..=1.0).contains(&pg) {
                return Err(BenchError::Config(
                    "--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -592,20 +594,20 @@ impl BenchConfig {
                    "--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
                ));
            }
-            if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
+            if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
                return Err(BenchError::Config(
                    "Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
                ));
            }
        }

-        if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
+        if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
            return Err(BenchError::Config(format!(
                "--steady-state-threshold must be in (0.0, 1.0], got {}",
-                cli.steady_state_threshold
+                args.steady_state_threshold
            )));
        }
-        if let Some(mw) = cli.steady_state_min_window
+        if let Some(mw) = args.steady_state_min_window
            && mw < 0.0
        {
            return Err(BenchError::Config(format!(
@@ -613,122 +615,122 @@ impl BenchConfig {
            )));
        }

-        if cli.profile_batch_threshold.is_some() && !cli.profile {
+        if args.profile_batch_threshold.is_some() && !args.profile {
            return Err(BenchError::Config(
                "--profile-batch-threshold requires --profile".into(),
            ));
        }
-        if cli.profile_duration <= 0.0 {
+        if args.profile_duration <= 0.0 {
            return Err(BenchError::Config(
                "--profile-duration must be positive".into(),
            ));
        }
-        if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
+        if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
            return Err(BenchError::Config(
                "--profile-duration requires --profile-batch-threshold".into(),
            ));
        }

        Ok(BenchConfig {
-            backend: cli.backend,
+            backend: args.backend,
            base_url,
            api_url,
-            model: cli.model.clone(),
-            model_name: cli.served_model_name.clone(),
+            model: args.model.clone(),
+            model_name: args.served_model_name.clone(),
            tokenizer_id,
-            tokenizer_mode: cli.tokenizer_mode.clone(),
-            trust_remote_code: cli.trust_remote_code,
-            skip_tokenizer_init: cli.skip_tokenizer_init,
-            dataset_name: cli.dataset_name,
-            dataset_path: cli.dataset_path.clone(),
-            max_model_len: cli.max_model_len,
+            tokenizer_mode: args.tokenizer_mode.clone(),
+            trust_remote_code: args.trust_remote_code,
+            skip_tokenizer_init: args.skip_tokenizer_init,
+            dataset_name: args.dataset_name,
+            dataset_path: args.dataset_path.clone(),
+            max_model_len: args.max_model_len,
            random_input_len,
            random_output_len,
-            random_prefix_len: cli.random_prefix_len,
+            random_prefix_len: args.random_prefix_len,
            random_range_ratio,
-            random_batch_size: cli.random_batch_size,
+            random_batch_size: args.random_batch_size,
            is_reranker,
-            custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
-            prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
-            prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
-            prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
-            prefix_repetition_output_len: cli
+            custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
+            prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
+            prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
+            prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
+            prefix_repetition_output_len: args
                .output_len
-                .unwrap_or(cli.prefix_repetition_output_len),
-            random_cache_hit_fraction: cli.random_cache_hit_fraction,
-            random_cache_ratio: cli.random_cache_ratio,
-            sharegpt_output_len: cli.sharegpt_output_len,
-            sonnet_input_len: cli.sonnet_input_len,
-            sonnet_output_len: cli.sonnet_output_len,
-            sonnet_prefix_len: cli.sonnet_prefix_len,
-            no_oversample: cli.no_oversample,
-            disable_shuffle: cli.disable_shuffle,
-            num_prompts: cli.num_prompts,
-            request_rate: cli.request_rate,
-            burstiness: cli.burstiness,
-            max_concurrency: cli.max_concurrency,
-            steady_state_threshold: cli.steady_state_threshold,
-            steady_state_min_window: cli.steady_state_min_window,
-            no_steady_state: cli.no_steady_state,
-            disable_tqdm: cli.disable_tqdm,
-            num_warmups: cli.num_warmups,
-            profile: cli.profile,
-            profile_batch_threshold: cli.profile_batch_threshold,
-            profile_duration: cli.profile_duration,
-            save_result: cli.save_result,
-            save_detailed: cli.save_detailed,
-            append_result: cli.append_result,
-            result_dir: cli.result_dir.clone(),
-            result_filename: cli.result_filename.clone(),
-            seed: cli.seed,
+                .unwrap_or(args.prefix_repetition_output_len),
+            random_cache_hit_fraction: args.random_cache_hit_fraction,
+            random_cache_ratio: args.random_cache_ratio,
+            sharegpt_output_len: args.sharegpt_output_len,
+            sonnet_input_len: args.sonnet_input_len,
+            sonnet_output_len: args.sonnet_output_len,
+            sonnet_prefix_len: args.sonnet_prefix_len,
+            no_oversample: args.no_oversample,
+            disable_shuffle: args.disable_shuffle,
+            num_prompts: args.num_prompts,
+            request_rate: args.request_rate,
+            burstiness: args.burstiness,
+            max_concurrency: args.max_concurrency,
+            steady_state_threshold: args.steady_state_threshold,
+            steady_state_min_window: args.steady_state_min_window,
+            no_steady_state: args.no_steady_state,
+            disable_tqdm: args.disable_tqdm,
+            num_warmups: args.num_warmups,
+            profile: args.profile,
+            profile_batch_threshold: args.profile_batch_threshold,
+            profile_duration: args.profile_duration,
+            save_result: args.save_result,
+            save_detailed: args.save_detailed,
+            append_result: args.append_result,
+            result_dir: args.result_dir.clone(),
+            result_filename: args.result_filename.clone(),
+            seed: args.seed,
            ignore_eos,
-            insecure: cli.insecure,
+            insecure: args.insecure,
            selected_percentile_metrics,
            selected_percentiles,
            sweep_summary_percentiles,
-            label: cli.label.clone(),
-            logprobs: cli.logprobs,
-            request_id_prefix: cli.get_request_id_prefix(),
-            ready_check_timeout_sec: cli.ready_check_timeout_sec,
+            label: args.label.clone(),
+            logprobs: args.logprobs,
+            request_id_prefix: args.get_request_id_prefix(),
+            ready_check_timeout_sec: args.ready_check_timeout_sec,
            extra_headers,
            extra_body,
            metadata,
-            dry_run: cli.dry_run,
+            dry_run: args.dry_run,
            goodput,
            ramp_up,
-            multi_turn: cli.multi_turn,
-            multi_turn_num_turns: cli.multi_turn_num_turns,
+            multi_turn: args.multi_turn,
+            multi_turn_num_turns: args.multi_turn_num_turns,
            multi_turn_min_turns,
            multi_turn_max_turns,
-            sharegpt_multi_turn_max_turns: if cli.multi_turn
-                && cli.dataset_name == DatasetName::ShareGpt
-                && cli.multi_turn_max_turns != 0
+            sharegpt_multi_turn_max_turns: if args.multi_turn
+                && args.dataset_name == DatasetName::ShareGpt
+                && args.multi_turn_max_turns != 0
            {
-                Some(cli.multi_turn_max_turns)
+                Some(args.multi_turn_max_turns)
            } else {
                None
            },
            per_turn_input_len,
-            multi_turn_concurrency: cli.multi_turn_concurrency,
-            multi_turn_delay_ms: cli.multi_turn_delay_ms,
-            multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
-            multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
-            speed_bench_config: cli.speed_bench_config,
-            speed_bench_category: cli.speed_bench_category.clone(),
-            speed_bench_max_input_len: cli.speed_bench_max_input_len,
-            hf_split: cli.hf_split.clone(),
-            hf_subset: cli.hf_subset.clone(),
-            hf_output_len: cli.hf_output_len,
-            hf_text_column: cli.hf_text_column.clone(),
-            reset_prefix_cache: cli.reset_prefix_cache,
-            prompt_token_ids: cli.prompt_token_ids,
-            random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
-            random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
+            multi_turn_concurrency: args.multi_turn_concurrency,
+            multi_turn_delay_ms: args.multi_turn_delay_ms,
+            multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
+            multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
+            speed_bench_config: args.speed_bench_config,
+            speed_bench_category: args.speed_bench_category.clone(),
+            speed_bench_max_input_len: args.speed_bench_max_input_len,
+            hf_split: args.hf_split.clone(),
+            hf_subset: args.hf_subset.clone(),
+            hf_output_len: args.hf_output_len,
+            hf_text_column: args.hf_text_column.clone(),
+            reset_prefix_cache: args.reset_prefix_cache,
+            prompt_token_ids: args.prompt_token_ids,
+            random_mm_base_items_per_request: args.random_mm_base_items_per_request,
+            random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
            random_mm_limit,
            random_mm_buckets,
-            enable_multimodal_chat: cli.enable_multimodal_chat,
+            enable_multimodal_chat: args.enable_multimodal_chat,
            lora_modules,
-            lora_assignment: cli.lora_assignment,
+            lora_assignment: args.lora_assignment,
        })
    }
 }
@@ -811,17 +813,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
    Ok(config)
 }

-fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
-    let strategy = match cli.ramp_up_strategy {
+fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
+    let strategy = match args.ramp_up_strategy {
        None => return Ok(None),
        Some(s) => s,
    };

-    let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
+    let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
        BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
    })?;

-    let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
+    let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
        BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
    })?;

@@ -843,7 +845,21 @@ mod tests {
    use clap::Parser;

    use super::*;
-    use crate::cli::Cli;
+    use crate::cli::BenchServeArgs;
+
+    #[derive(Parser)]
+    struct TestCli {
+        #[command(flatten)]
+        args: BenchServeArgs,
+    }
+
+    fn parse_args<I, T>(args: I) -> BenchServeArgs
+    where
+        I: IntoIterator<Item = T>,
+        T: Into<std::ffi::OsString> + Clone,
+    {
+        TestCli::parse_from(args).args
+    }

    fn base_multi_turn_args() -> Vec<&'static str> {
        vec![
@@ -859,8 +875,8 @@ mod tests {
    #[test]
    fn test_prefix_sharing_defaults_to_zero() {
        let args = base_multi_turn_args();
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
        assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
        assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
    }
@@ -874,8 +890,8 @@ mod tests {
            "--multi-turn-prefix-conversation-ratio",
            "0.8",
        ]);
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
        assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
        assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
    }
@@ -889,8 +905,8 @@ mod tests {
            "--multi-turn-prefix-conversation-ratio",
            "0.6",
        ]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
    }

    #[test]
@@ -902,16 +918,16 @@ mod tests {
            "--multi-turn-prefix-conversation-ratio",
            "0.5",
        ]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
    }

    #[test]
    fn test_prefix_sharing_out_of_range_fails() {
        let mut args = base_multi_turn_args();
        args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
    }

    #[test]
@@ -928,8 +944,8 @@ mod tests {
            "--multi-turn-prefix-global-ratio",
            "0.1",
        ];
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
    }

    #[test]
@@ -944,8 +960,8 @@ mod tests {
            "--dataset-name",
            "sharegpt",
        ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();

        assert_eq!(config.multi_turn_max_turns, 3);
        assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -968,8 +984,8 @@ mod tests {
            "--multi-turn-max-turns",
            "2",
        ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();
        assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
    }

@@ -987,8 +1003,8 @@ mod tests {
            "--multi-turn-max-turns",
            "1",
        ];
-        let cli = Cli::parse_from(args);
-        let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
+        let args = parse_args(args);
+        let err = BenchConfig::from_args(&args).unwrap_err().to_string();
        assert!(
            err.contains("at least 2 for ShareGPT"),
            "expected ShareGPT-specific error, got: {err}"
@@ -1009,8 +1025,8 @@ mod tests {
            "--multi-turn-max-turns",
            "20",
        ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();

        assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
    }
@@ -1018,8 +1034,8 @@ mod tests {
    #[test]
    fn test_sweep_summary_percentiles_default_empty() {
        let args = base_multi_turn_args();
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();

        assert!(config.sweep_summary_percentiles.is_empty());
        assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1034,8 +1050,8 @@ mod tests {
            "--sweep-summary-percentiles",
            "90,95,90",
        ]);
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();

        assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
        assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1045,8 +1061,8 @@ mod tests {
    fn test_invalid_sweep_summary_percentile_fails() {
        let mut args = base_multi_turn_args();
        args.extend(["--sweep-summary-percentiles", "101"]);
-        let cli = Cli::parse_from(args);
-        assert!(BenchConfig::from_cli(&cli).is_err());
+        let args = parse_args(args);
+        assert!(BenchConfig::from_args(&args).is_err());
    }

    #[test]
@@ -1058,12 +1074,20 @@ mod tests {
            "--max-model-len",
            "4096",
        ];
-        let cli = Cli::parse_from(args);
-        let config = BenchConfig::from_cli(&cli).unwrap();
+        let args = parse_args(args);
+        let config = BenchConfig::from_args(&args).unwrap();

        assert_eq!(config.max_model_len, Some(4096));
    }

+    #[test]
+    fn test_tokenizer_id_deferred_when_model_is_unspecified() {
+        let args = parse_args(["vllm-bench"]);
+        let config = BenchConfig::from_args(&args).unwrap();
+
+        assert_eq!(config.tokenizer_id, None);
+    }
+
    #[test]
    fn test_zero_max_model_len_fails() {
        let args = vec![
@@ -1073,9 +1097,9 @@ mod tests {
            "--max-model-len",
            "0",
        ];
-        let cli = Cli::parse_from(args);
+        let args = parse_args(args);

-        assert!(BenchConfig::from_cli(&cli).is_err());
+        assert!(BenchConfig::from_args(&args).is_err());
    }
    #[test]
    fn test_range_ratio_parse_float() {
@@ -40,8 +40,11 @@ impl HubRepo {
                .build()
                .map_err(|e| format!("Failed to build download runtime: {e}"))?;
            rt.block_on(async move {
-                let api = hf_hub::api::tokio::Api::new()
-                    .map_err(|e| format!("Failed to init HF API: {e}"))?;
+                let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
+                if let Ok(token) = std::env::var("HF_TOKEN") {
+                    builder = builder.with_token(Some(token));
+                }
+                let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
                api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
            })
        })
@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+mod backends;
+mod benchmark;
+mod cli;
+mod compare;
+mod config;
+mod datasets;
+mod error;
+mod hub;
+mod metrics;
+mod multi_run;
+mod multi_turn;
+mod output;
+mod rate_control;
+mod ready_checker;
+mod sweep;
+mod tiktoken;
+mod tokenizer;
+
+use anyhow::Context;
+
+pub use cli::{
+    BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
+};
+use config::BenchConfig;
+
+/// Prepare process-wide resources for a benchmark run.
+pub fn prepare_process() {
+    // Raise the open-file soft limit to the hard limit. High-concurrency
+    // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
+    if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
+        && new > 1024
+    {
+        eprintln!("Open-file limit: {new}");
+    }
+}
+
+/// Run the online serving benchmark.
+pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
+    // --- Compare mode: no server needed, just diff two JSON files ---
+    if let Some(ref files) = args.compare {
+        return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
+    }
+
+    let config = BenchConfig::from_args(&args).context("Configuration error")?;
+
+    async {
+        if config.multi_turn {
+            if let Some(ref sweep_mc) = args.sweep_max_concurrency {
+                // --- Sweep over concurrency in multi-turn mode ---
+                let values = sweep::parse_concurrency_values(sweep_mc)
+                    .context("Invalid --sweep-max-concurrency")?;
+                sweep::run_multi_turn_concurrency_sweep(
+                    &config,
+                    &values,
+                    args.sweep_num_prompts_factor,
+                )
+                .await?;
+            } else {
+                // --- Single multi-turn conversation benchmark ---
+                multi_turn::run_multi_turn_benchmark(&config).await?;
+            }
+        } else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
+            // --- Sweep over max-concurrency ---
+            let values = sweep::parse_concurrency_values(sweep_mc)
+                .context("Invalid --sweep-max-concurrency")?;
+            sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
+        } else if let Some(ref sweep_rate) = args.sweep_request_rate {
+            // --- Sweep over request-rate ---
+            let values =
+                sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
+            sweep::run_rate_sweep(&config, &values).await?;
+        } else if args.num_runs > 1 {
+            // --- Multi-run with statistical aggregation ---
+            multi_run::run_multi(&config, args.num_runs).await?;
+        } else {
+            // --- Normal single benchmark ---
+            benchmark::run_benchmark(&config).await?;
+        }
+        anyhow::Ok(())
+    }
+    .await
+    .context("Benchmark failed")
+}
@@ -1,92 +1,32 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

-mod backends;
-mod benchmark;
-mod cli;
-mod compare;
-mod config;
-mod datasets;
-mod error;
-mod hub;
-mod metrics;
-mod multi_run;
-mod multi_turn;
-mod output;
-mod rate_control;
-mod ready_checker;
-mod sweep;
-mod tiktoken;
-mod tokenizer;
-
 #[cfg(not(target_env = "msvc"))]
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;

 use anyhow::Context;
 use clap::Parser;
-use cli::Cli;
-use config::BenchConfig;
+
+#[derive(Parser)]
+#[command(
+    name = "vllm-bench",
+    about = "Benchmark online serving throughput",
+    version
+)]
+struct Cli {
+    #[command(flatten)]
+    args: vllm_bench::BenchServeArgs,
+}

 fn main() -> anyhow::Result<()> {
-    // Raise the open-file soft limit to the hard limit. High-concurrency
-    // benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
-    if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
-        && new > 1024
-    {
-        eprintln!("Open-file limit: {new}");
-    }
-
    let cli = Cli::parse();
-
-    // --- Compare mode: no server needed, just diff two JSON files ---
-    if let Some(ref files) = cli.compare {
-        return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
-    }
-
-    let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
+    vllm_bench::prepare_process();

    let runtime = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
-        .expect("Failed to build tokio runtime");
+        .context("Failed to build tokio runtime")?;

-    runtime
-        .block_on(async {
-            if config.multi_turn {
-                if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
-                    // --- Sweep over concurrency in multi-turn mode ---
-                    let values = sweep::parse_concurrency_values(sweep_mc)
-                        .context("Invalid --sweep-max-concurrency")?;
-                    sweep::run_multi_turn_concurrency_sweep(
-                        &config,
-                        &values,
-                        cli.sweep_num_prompts_factor,
-                    )
-                    .await?;
-                } else {
-                    // --- Single multi-turn conversation benchmark ---
-                    multi_turn::run_multi_turn_benchmark(&config).await?;
-                }
-            } else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
-                // --- Sweep over max-concurrency ---
-                let values = sweep::parse_concurrency_values(sweep_mc)
-                    .context("Invalid --sweep-max-concurrency")?;
-                sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
-                    .await?;
-            } else if let Some(ref sweep_rate) = cli.sweep_request_rate {
-                // --- Sweep over request-rate ---
-                let values =
-                    sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
-                sweep::run_rate_sweep(&config, &values).await?;
-            } else if cli.num_runs > 1 {
-                // --- Multi-run with statistical aggregation ---
-                multi_run::run_multi(&config, cli.num_runs).await?;
-            } else {
-                // --- Normal single benchmark ---
-                benchmark::run_benchmark(&config).await?;
-            }
-            anyhow::Ok(())
-        })
-        .context("Benchmark failed")
+    runtime.block_on(vllm_bench::run(cli.args))
 }
@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
            let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
            let (value, field) = match spec.field_layout_for(key) {
                Some(FieldLayout::Batched) => (
-                    tensor.batched_value_at(index)?,
+                    tensor.batched_wire_value_at(index)?,
                    MmField::Batched(MmBatchedField { keep_on_cpu }),
                ),
                Some(FieldLayout::Flat { sizes_key }) => {
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
                    })?;
                    let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
                    (
-                        tensor.flat_value_range(start, end)?,
+                        tensor.flat_wire_value_range(start, end)?,
                        MmField::Flat(MmFlatField {
                            slices: vec![MmSlice::Slice(SliceSpec {
                                start: Some(0),
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
                    )
                }
                None => (
-                    tensor.clone(),
+                    tensor.try_into()?,
                    MmField::Shared(MmSharedField {
                        batch_size: len,
                        keep_on_cpu,
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
            data.insert(
                key.clone(),
                MmFieldElem {
-                    data: Some(value.try_into()?),
+                    data: Some(value),
                    field,
                },
            );
@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
 use crate::error::{Error, Result, bail_multimodal, multimodal};

 /// Representation for multimodal kwarg values for transformation.
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub(super) enum KwargValue {
    /// Float tensor with row-major flat data and shape.
    F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,28 +107,19 @@ impl KwargValue {
    }
 }

-impl TryFrom<KwargValue> for ProtocolKwargValue {
+impl TryFrom<&KwargValue> for ProtocolKwargValue {
    type Error = Error;

-    fn try_from(value: KwargValue) -> Result<Self> {
-        match value {
-            KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
-                WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
-            )),
-            KwargValue::Passthrough(value) => Ok(value),
-        }
+    fn try_from(value: &KwargValue) -> Result<Self> {
+        let tensor = match value {
+            KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
+            KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
+            KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
+            KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
+            KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
+            KwargValue::Passthrough(value) => return Ok(value.clone()),
+        };
+        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
    }
 }

@@ -145,63 +136,55 @@ impl KwargValue {
        }
    }

-    /// Extract one media item from a batched tensor field.
+    /// Convert one media item from a batched tensor field to wire bytes.
    ///
    /// Batched fields use their first axis as media-item index and drop that
    /// axis in the per-feature value, matching vLLM's batched-field semantics.
-    pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
-        match self {
-            Self::F32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::F32Tensor { data, shape })
-            }
-            Self::F16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::F16Tensor { data, shape })
-            }
-            Self::Bf16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::Bf16Tensor { data, shape })
-            }
-            Self::I64Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::I64Tensor { data, shape })
-            }
-            Self::U32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
-                Ok(Self::U32Tensor { data, shape })
-            }
-            Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
-        }
+    pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
+        self.wire_value_range(index, index + 1, true)
    }

-    /// Extract one media item's variable-length range from a flat tensor field.
+    /// Convert one media item's flat tensor range directly to wire bytes.
    ///
    /// Flat fields keep the first axis as the sliced length for this item.
-    pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
-        match self {
+    pub(super) fn flat_wire_value_range(
+        &self,
+        start: usize,
+        end: usize,
+    ) -> Result<ProtocolKwargValue> {
+        self.wire_value_range(start, end, false)
+    }
+
+    fn wire_value_range(
+        &self,
+        start: usize,
+        end: usize,
+        drop_axis: bool,
+    ) -> Result<ProtocolKwargValue> {
+        let tensor = match self {
            Self::F32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::F32Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_f32(shape, data)
            }
            Self::F16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::F16Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_f16(shape, data)
            }
            Self::Bf16Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::Bf16Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_bf16(shape, data)
            }
            Self::I64Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::I64Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_i64(shape, data)
            }
            Self::U32Tensor { data, shape } => {
-                let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
-                Ok(Self::U32Tensor { data, shape })
+                let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
+                WireTensor::from_u32(shape, data)
            }
-            Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
-        }
+            Self::Passthrough(value) => return Ok(value.clone()),
+        };
+        tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
    }
 }

@@ -240,13 +223,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
 }

 /// Slice a flat row-major tensor along its first axis.
-fn slice_first_axis_range<T: Clone>(
+fn slice_first_axis_range<'a, T>(
    shape: &[usize],
-    data: &[T],
+    data: &'a [T],
    start: usize,
    end: usize,
    drop_axis: bool,
-) -> Result<(Vec<usize>, Vec<T>)> {
+) -> Result<(Vec<usize>, &'a [T])> {
    let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
    if start > end || end > first_dim {
        bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -270,7 +253,7 @@ fn slice_first_axis_range<T: Clone>(
        shape[0] = end - start;
        shape
    };
-    Ok((out_shape, data[data_start..data_end].to_vec()))
+    Ok((out_shape, &data[data_start..data_end]))
 }

 #[cfg(test)]
@@ -278,35 +261,39 @@ mod tests {
    use super::*;

    #[test]
-    fn batched_value_at_drops_first_axis() {
+    fn batched_wire_value_at_drops_first_axis() {
        let value = KwargValue::F32Tensor {
            data: vec![1.0, 2.0, 3.0, 4.0],
            shape: vec![2, 2],
        };

-        let value = value.batched_value_at(1).unwrap();
+        let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
+            panic!("expected tensor");
+        };

-        assert!(matches!(
-            value,
-            KwargValue::F32Tensor { data, shape }
-                if shape == vec![2] && data == vec![3.0, 4.0]
-        ));
+        assert_eq!(tensor.shape, vec![2]);
+        assert_eq!(
+            tensor.data.into_raw_view().unwrap(),
+            [3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
+        );
    }

    #[test]
-    fn flat_value_range_keeps_first_axis() {
+    fn flat_wire_value_range_keeps_first_axis() {
        let value = KwargValue::U32Tensor {
            data: (0..10).collect(),
            shape: vec![5, 2],
        };

-        let value = value.flat_value_range(1, 3).unwrap();
+        let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
+            panic!("expected tensor");
+        };

-        assert!(matches!(
-            value,
-            KwargValue::U32Tensor { data, shape }
-                if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
-        ));
+        assert_eq!(tensor.shape, vec![2, 2]);
+        assert_eq!(
+            tensor.data.into_raw_view().unwrap(),
+            [2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
+        );
    }

    #[test]
@@ -336,7 +323,7 @@ mod tests {
        let value =
            KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();

-        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
+        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
        else {
            panic!("expected tensor");
        };
@@ -351,7 +338,7 @@ mod tests {
        let value =
            KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();

-        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
+        let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
        else {
            panic!("expected tensor");
        };
@@ -130,7 +130,7 @@ fn build_video_item(
        let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
        let (value, field) = match support.spec.field_layout_for(&key) {
            Some(FieldLayout::Batched) => (
-                tensor.batched_value_at(0)?,
+                tensor.batched_wire_value_at(0)?,
                MmField::Batched(MmBatchedField { keep_on_cpu }),
            ),
            Some(FieldLayout::Flat { .. }) => {
@@ -138,7 +138,7 @@ fn build_video_item(
                    .first_dim()
                    .ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
                (
-                    tensor,
+                    (&tensor).try_into()?,
                    MmField::Flat(MmFlatField {
                        slices: vec![MmSlice::Slice(SliceSpec {
                            start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
                )
            }
            None => (
-                tensor,
+                (&tensor).try_into()?,
                MmField::Shared(MmSharedField {
                    batch_size: 1,
                    keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
        data.insert(
            key,
            MmFieldElem {
-                data: Some(value.try_into()?),
+                data: Some(value),
                field,
            },
        );
@@ -74,7 +74,7 @@ impl DefaultChatOutputProcessor {
            Box::new(CombinedParser::new(reasoning_parser, tool_parser)) as Box<dyn UnifiedParser>
        };

-        apply_structural_tag_constraint(request, parser.structural_tag_model())?;
+        apply_structural_tag_constraint(request, parser.structural_tag_builder())?;

        if parser.preserve_special_tokens() {
            request.decode_options.skip_special_tokens = false;
@@ -7,7 +7,8 @@ use thiserror_ext::AsReport;
 use vllm_engine_core_client::protocol::structured_outputs::{
    StructuredOutputBackend, StructuredOutputsParams,
 };
-use vllm_parser::tool::StructuralTagModel;
+use vllm_parser::tool::StructuralTagBuilder;
+use xgrammar_structural_tag::builders::StructuralTagOptions;
 use xgrammar_structural_tag::{
    FunctionDefinition, FunctionToolParam, ToolChoice as StructuralTagToolChoice, ToolParam,
    build_structural_tag,
@@ -20,9 +21,9 @@ use crate::{Error, Result as ChatResult};
 /// support and the request's tool choice.
 pub(super) fn apply_structural_tag_constraint(
    request: &mut ChatRequest,
-    model: Option<StructuralTagModel>,
+    builder: Option<&dyn StructuralTagBuilder>,
 ) -> ChatResult<()> {
-    let Some(model) = model else {
+    let Some(builder) = builder else {
        return Ok(());
    };
    let Some(tool_choice) = structural_tag_tool_choice(request) else {
@@ -42,11 +43,16 @@ pub(super) fn apply_structural_tag_constraint(
        })
        .collect::<Vec<_>>();

-    let structural_tag = build_structural_tag(model, &tools, tool_choice, false)
-        .and_then(|tag| tag.to_json_string())
-        .map_err(|error| Error::StructuralTag {
-            message: error.to_report_string(),
-        })?;
+    let structural_tag = build_structural_tag(
+        builder,
+        &tools,
+        tool_choice,
+        StructuralTagOptions::default().with_reasoning(false),
+    )
+    .and_then(|tag| tag.to_json_string())
+    .map_err(|error| Error::StructuralTag {
+        message: error.to_report_string(),
+    })?;

    // Overwrite any existing structured output settings with the structural tag constraint.
    request.sampling_params.structured_outputs = Some(StructuredOutputsParams {
@@ -141,7 +147,7 @@ mod tests {
        let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]);
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag should build");

        let tag = structural_tag_value(&request);
@@ -154,7 +160,7 @@ mod tests {
        let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]);
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag decision should succeed");

        assert!(request.sampling_params.structured_outputs.is_none());
@@ -169,7 +175,7 @@ mod tests {
        });
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag should build");

        let params = structured_outputs(&request);
@@ -184,7 +190,7 @@ mod tests {
        let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]);
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag should build");

        let tag = structural_tag_value(&request);
@@ -201,7 +207,7 @@ mod tests {
        });
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag should build");

        let params = structured_outputs(&request);
@@ -221,7 +227,7 @@ mod tests {
        );
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag should build");

        let tag = structural_tag_value(&request).to_string();
@@ -234,7 +240,7 @@ mod tests {
        let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]);
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag decision should succeed");

        assert!(request.sampling_params.structured_outputs.is_none());
@@ -249,7 +255,7 @@ mod tests {
        });
        let parser = qwen3_coder_parser(&request.tools);

-        apply_structural_tag_constraint(&mut request, parser.structural_tag_model())
+        apply_structural_tag_constraint(&mut request, parser.structural_tag_builder())
            .expect("structural tag decision should succeed");

        let params = structured_outputs(&request);
@@ -236,9 +236,10 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {

    loops.into_iter().any(|loop_ast| {
        matches!(loop_ast.target, Expr::Var(_))
-            && message_varnames
-                .iter()
-                .any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
+            && (is_var_access(&loop_ast.iter, "content")
+                || message_varnames.iter().any(|varname| {
+                    is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
+                }))
    })
 }

@@ -315,6 +316,16 @@ mod tests {
        );
    }

+    #[test]
+    fn detects_openai_template_with_content_parameter_loop() {
+        assert_eq!(
+            detect(
+                "{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}"
+            ),
+            ChatTemplateContentFormat::OpenAi
+        );
+    }
+
    #[test]
    fn detects_openai_template_with_messages_alias() {
        assert_eq!(
@@ -1309,6 +1309,26 @@ mod tests {
        .assert_eq(&rendered);
    }

+    #[test]
+    fn qwen35_template_auto_detects_openai_multimodal_content() {
+        let mut request = image_request();
+        request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
+
+        let rendered = render_mm(
+            QWEN3_5_0_8B_TEMPLATE,
+            &request,
+            ChatTemplateContentFormatOption::Auto,
+        )
+        .unwrap();
+
+        expect![[r#"
+            Text(
+                "<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n",
+            )
+        "#]]
+        .assert_debug_eq(&rendered.prompt);
+    }
+
    #[test]
    fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
        let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
@@ -29,6 +29,7 @@ tokio-util.workspace = true
 tracing.workspace = true
 tracing-subscriber.workspace = true
 uuid.workspace = true
+vllm-bench.workspace = true
 vllm-chat.workspace = true
 vllm-engine-core-client.workspace = true
 vllm-managed-engine.workspace = true
@@ -79,13 +79,23 @@ impl Cli {
 }

 /// Supported top-level CLI commands.
-#[derive(Debug, Subcommand, PartialEq, Eq)]
+#[derive(Debug, Subcommand)]
 pub enum Command {
    /// Run the Rust OpenAI frontend as a Python-supervised worker.
    Frontend(FrontendArgs),
    /// Launch a managed Python headless engine, then run the Rust OpenAI
    /// frontend.
    Serve(ServeArgs),
+    /// Run vLLM benchmarks.
+    #[command(subcommand)]
+    Bench(BenchCommand),
+}
+
+/// Supported benchmark commands.
+#[derive(Debug, Subcommand)]
+pub enum BenchCommand {
+    /// Benchmark online serving throughput.
+    Serve(vllm_bench::BenchServeArgs),
 }

 /// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
@@ -5,7 +5,27 @@ use expect_test::expect;
 use vllm_engine_core_client::TransportMode;
 use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};

-use super::{Cli, Command};
+use super::{BenchCommand, Cli, Command};
+
+#[test]
+fn bench_serve_args_parse_without_managed_engine_repartition() {
+    let cli = Cli::try_parse_from([
+        "vllm-rs",
+        "bench",
+        "serve",
+        "--backend",
+        "openai-chat",
+        "--request-rate",
+        "inf",
+    ])
+    .unwrap();
+
+    let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
+        panic!("expected bench serve args");
+    };
+    assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
+    assert!(args.request_rate.is_infinite());
+}

 #[test]
 fn serve_args_forward_python_flags_with_separator() {
@@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken;
 use tracing::{info, warn};
 use vllm_managed_engine::ManagedEngineHandle;

-use crate::cli::{Cli, Command};
+use crate::cli::{BenchCommand, Cli, Command};

 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -100,6 +100,10 @@ fn main() -> Result<()> {
 async fn async_main(cli: Cli) -> Result<()> {
    match cli.command {
        Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
+        Command::Bench(BenchCommand::Serve(bench_args)) => {
+            vllm_bench::prepare_process();
+            vllm_bench::run(bench_args).await
+        }
        Command::Serve(args) => {
            let handshake_port = args.managed_engine.resolve_handshake_port()?;

@@ -55,52 +55,57 @@ pub struct WireNdArray {

 impl WireNdArray {
    /// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
+    pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
+        let data = data.as_ref();
        validate_element_count(&shape, data.len())?;
        Ok(Self {
            dtype: "float32".to_string(),
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
        })
    }

    /// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
+    pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
+        let data = data.as_ref();
        validate_element_count(&shape, data.len())?;
        Ok(Self {
            dtype: "float16".to_string(),
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
        })
    }

    /// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
+    pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
+        let data = data.as_ref();
        validate_element_count(&shape, data.len())?;
        Ok(Self {
            dtype: "bfloat16".to_string(),
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
        })
    }

    /// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
+    pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
+        let data = data.as_ref();
        validate_element_count(&shape, data.len())?;
        Ok(Self {
            dtype: "int64".to_string(),
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
        })
    }

    /// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
-    pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
+    pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
+        let data = data.as_ref();
        validate_element_count(&shape, data.len())?;
        Ok(Self {
            dtype: "uint32".to_string(),
            shape,
-            data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)),
+            data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
        })
    }

@@ -4,7 +4,7 @@
 use std::sync::Arc;

 use vllm_parser::tool::{
-    Result, StructuralTagModel, Tool, ToolParser, ToolParserError, ToolParserOutput,
+    Result, StructuralTagBuilder, Tool, ToolParser, ToolParserError, ToolParserOutput,
 };
 use vllm_parser::unified::{
    UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
@@ -85,8 +85,8 @@ impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
        self.inner.preserve_special_tokens()
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        self.inner.structural_tag_model()
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        self.inner.structural_tag_builder()
    }

    fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use super::{DeepSeekDsmlToolParser, DsmlTokens};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 /// Tool parser for DeepSeek V3.2 models.
 ///
@@ -47,8 +47,8 @@ impl ToolParser for DeepSeekV32ToolParser {
        true
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::DeepSeekV32)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::DeepSeekV32.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use super::{DeepSeekDsmlToolParser, DsmlTokens};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 /// Tool parser for DeepSeek V4 models.
 ///
@@ -50,8 +50,8 @@ impl ToolParser for DeepSeekV4ToolParser {
        true
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::DeepSeekV4)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::DeepSeekV4.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -73,7 +73,7 @@ mod tests {

    use super::DeepSeekV4ToolParser;
    use crate::tool::test_utils::{collect_stream, test_tools};
-    use crate::tool::{StructuralTagModel, ToolParser, ToolParserTestExt as _};
+    use crate::tool::{ToolParser, ToolParserTestExt as _};

    fn build_tool_call(function_name: &str, params: &[(&str, &str)]) -> String {
        let params = params
@@ -91,13 +91,10 @@ mod tests {
    }

    #[test]
-    fn deepseek_v4_exposes_structural_tag_model() {
+    fn deepseek_v4_exposes_structural_tag_builder() {
        let parser = DeepSeekV4ToolParser::new(&test_tools());

-        assert_eq!(
-            parser.structural_tag_model(),
-            Some(StructuralTagModel::DeepSeekV4)
-        );
+        assert!(parser.structural_tag_builder().is_some());
    }

    #[test]
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 /// Tool parser for DeepSeek V3 JSON-fenced tool calls.
 ///
@@ -35,8 +35,8 @@ impl ToolParser for DeepSeekV3ToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::DeepSeekR1)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::DeepSeekR1.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use super::{DeepSeekJsonFormat, DeepSeekJsonToolParser};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 /// Tool parser for DeepSeek V3.1 raw JSON tool calls.
 ///
@@ -31,8 +31,8 @@ impl ToolParser for DeepSeekV31ToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::DeepSeekV31)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::DeepSeekV31.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use super::{GlmXmlToolParser, Separator};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 /// Tool parser for GLM-4.7 MoE XML-style tool calls.
 ///
@@ -25,8 +25,8 @@ impl ToolParser for Glm47MoeToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::Glm47)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::Glm47.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
 use super::parameters::ToolSchemas;
 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::{StructuralTagModel, Tool};
+use crate::tool::{StructuralTagBuilder, Tool};

 const TOOL_CALLS_START: &str = "<tool_calls>";
 const TOOL_CALLS_END: &str = "</tool_calls>";
@@ -116,8 +116,8 @@ impl ToolParser for HyV3ToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::HyV3)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::HyV3.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 const HERMES_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
    parser_name: "Hermes",
@@ -48,8 +48,8 @@ impl ToolParser for HermesToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::Hermes)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::Hermes.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -12,7 +12,9 @@ use super::{
    argument_delta_event, tool_call_header_event,
 };
 use crate::tool::utils::{JsonObjectScanState, parse_buffered_event};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolCallDelta, ToolParser, ToolParserOutput};
+use crate::tool::{
+    Result, StructuralTagBuilder, Tool, ToolCallDelta, ToolParser, ToolParserOutput,
+};

 #[derive(Debug, Clone, PartialEq, Eq)]
 enum LlamaJsonMode {
@@ -136,8 +138,8 @@ impl ToolParser for Llama3JsonToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::Llama)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::Llama.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 use super::{JsonToolCallConfig, JsonToolCallParser, JsonToolCallWhitespace};
-use crate::tool::{Result, StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{Result, StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 const QWEN_XML_CONFIG: JsonToolCallConfig = JsonToolCallConfig {
    parser_name: "Qwen XML",
@@ -50,8 +50,8 @@ impl ToolParser for Qwen3XmlToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::Qwen3)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::Qwen3.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -11,7 +11,7 @@ use winnow::token::{literal, rest, take_until, take_while};

 use super::utils::{JsonObjectScanState, parse_buffered_event, safe_text_len, take_json_object};
 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::{StructuralTagModel, Tool};
+use crate::tool::{StructuralTagBuilder, Tool};

 const TOOL_CALLS_START: &str = "<|tool_calls_section_begin|>";
 const TOOL_CALLS_END: &str = "<|tool_calls_section_end|>";
@@ -150,8 +150,8 @@ impl ToolParser for KimiK2ToolParser {
        true
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::Kimi)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::Kimi.builder())
    }

    fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -10,7 +10,7 @@ use winnow::token::{literal, rest, take_until};
 use super::parameters::ToolSchemas;
 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
 use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::{StructuralTagModel, Tool};
+use crate::tool::{StructuralTagBuilder, Tool};

 const TOOL_CALL_START: &str = "<minimax:tool_call>";
 const TOOL_CALL_END: &str = "</minimax:tool_call>";
@@ -115,8 +115,8 @@ impl ToolParser for MinimaxM2ToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::Minimax)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::Minimax.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -36,7 +36,7 @@ pub use qwen_coder::Qwen3CoderToolParser;
 pub use seed_oss::SeedOssToolParser;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
-pub use xgrammar_structural_tag::Model as StructuralTagModel;
+pub use xgrammar_structural_tag::builders::StructuralTagBuilder;

 use crate::utils;

@@ -187,8 +187,8 @@ pub trait ToolParser: Send {
        false
    }

-    /// Return the xgrammar structural-tag model used for strict tool calling.
-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+    /// Return the xgrammar structural-tag builder used for strict tool calling.
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
        None
    }

@@ -9,8 +9,8 @@ use winnow::token::{literal, take_until};

 use super::parameters::ToolSchemas;
 use super::utils::{MarkerScanState, parse_buffered_event, safe_text_len, take_until_marker};
-use super::{Result, StructuralTagModel, ToolCallDelta, ToolParser, ToolParserOutput};
-use crate::tool::Tool;
+use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
+use crate::tool::{StructuralTagBuilder, Tool};

 const TOOL_CALL_START: &str = "<tool_call>";
 const TOOL_CALL_END: &str = "</tool_call>";
@@ -146,8 +146,8 @@ impl ToolParser for Qwen3CoderToolParser {
        Ok(Box::new(Self::new(tools)))
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        Some(StructuralTagModel::Qwen3Coder)
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        Some(xgrammar_structural_tag::Model::Qwen3Coder.builder())
    }

    fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
@@ -294,7 +294,7 @@ mod tests {
    use serde_json::{Value, json};
    use thiserror_ext::AsReport;

-    use super::{Qwen3CoderToolParser, StructuralTagModel, ToolParser};
+    use super::{Qwen3CoderToolParser, ToolParser};
    use crate::tool::test_utils::{collect_stream, split_by_chars, test_tools};
    use crate::tool::{ToolParserOutput, ToolParserTestExt as _};

@@ -308,13 +308,10 @@ mod tests {
    }

    #[test]
-    fn qwen_coder_exposes_structural_tag_model() {
+    fn qwen_coder_exposes_structural_tag_builder() {
        let parser = Qwen3CoderToolParser::new(&test_tools());

-        assert_eq!(
-            parser.structural_tag_model(),
-            Some(StructuralTagModel::Qwen3Coder)
-        );
+        assert!(parser.structural_tag_builder().is_some());
    }

    #[test]
@@ -7,7 +7,7 @@ use vllm_tokenizer::DynTokenizer;

 use super::{Result, UnifiedParser, UnifiedParserError, UnifiedParserOutput};
 use crate::reasoning::ReasoningParser;
-use crate::tool::{StructuralTagModel, Tool, ToolParser, ToolParserOutput};
+use crate::tool::{StructuralTagBuilder, Tool, ToolParser, ToolParserOutput};

 /// Unified parser that composes existing reasoning and tool parsers.
 pub struct CombinedParser {
@@ -79,8 +79,8 @@ impl UnifiedParser for CombinedParser {
            || self.tool.as_ref().is_some_and(|parser| parser.preserve_special_tokens())
    }

-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
-        self.tool.as_ref().and_then(|parser| parser.structural_tag_model())
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
+        self.tool.as_ref().and_then(|parser| parser.structural_tag_builder())
    }

    fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
@@ -269,10 +269,7 @@ mod tests {
    fn combined_parser_emits_tool_calls_from_visible_content() {
        let tool = Qwen3XmlToolParser::create(&test_tools()).unwrap();
        let mut parser = CombinedParser::new(None, Some(tool));
-        assert!(matches!(
-            parser.structural_tag_model(),
-            Some(crate::tool::StructuralTagModel::Qwen3)
-        ));
+        assert!(parser.structural_tag_builder().is_some());

        let output = collect(
            &mut parser,
@@ -16,7 +16,7 @@ use vllm_tokenizer::DynTokenizer;

 use crate::reasoning::ReasoningError;
 use crate::tool::{
-    StructuralTagModel, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
+    StructuralTagBuilder, Tool, ToolCallDelta, ToolParserError, ToolParserEvent, ToolParserOutput,
 };

 /// Result alias for unified parser operations.
@@ -171,8 +171,8 @@ pub trait UnifiedParser: Send {
        false
    }

-    /// Return the xgrammar structural-tag model used for strict tool calling.
-    fn structural_tag_model(&self) -> Option<StructuralTagModel> {
+    /// Return the xgrammar structural-tag builder used for strict tool calling.
+    fn structural_tag_builder(&self) -> Option<&dyn StructuralTagBuilder> {
        None
    }

@@ -238,13 +238,18 @@ fn collect_generate(
        None
    };
    let prompt_logprobs = if include_prompt_logprobs {
-        let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
-            ApiError::server_error(
-                "raw generate response requested prompt_logprobs but generation returned none"
-                    .to_string(),
-            )
-        })?;
-        Some(raw_prompt_logprobs_to_maps(prompt_logprobs))
+        match collected.prompt_logprobs.as_ref() {
+            Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)),
+            // A single-token prompt has no scored positions; same mapping
+            // as /v1/completions.
+            None if collected.prompt_token_ids.len() == 1 => Some(vec![None]),
+            None => {
+                return Err(ApiError::server_error(
+                    "raw generate response requested prompt_logprobs but generation returned none"
+                        .to_string(),
+                ));
+            }
+        }
    } else {
        None
    };
@@ -472,4 +477,48 @@ mod tests {
            Some(2)
        );
    }
+
+    #[test]
+    fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
+        let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
+            request_id: "raw-1".to_string(),
+            prompt_logprobs: None,
+            token_ids: vec![3],
+            logprobs: None,
+            finish_reason: FinishReason::stop_eos(),
+            usage: vllm_llm::TokenUsage {
+                prompt_token_count: prompt_token_ids.len(),
+                output_token_count: 1,
+                cached_token_count: 0,
+            },
+            kv_transfer_params: None,
+            ec_transfer_params: None,
+            prompt_token_ids,
+        };
+
+        let response = collect_generate(
+            output_without_payload(vec![9707]),
+            "raw-1".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions {
+                include_prompt_logprobs: true,
+                ..Default::default()
+            },
+        )
+        .expect("single-token prompt without payload maps to [None]");
+        let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
+        assert_eq!(prompt_logprobs.len(), 1);
+        assert!(prompt_logprobs[0].is_none());
+
+        collect_generate(
+            output_without_payload(vec![9707, 11]),
+            "raw-2".to_string(),
+            ApiServerOptions::default(),
+            ResponseOptions {
+                include_prompt_logprobs: true,
+                ..Default::default()
+            },
+        )
+        .expect_err("multi-token prompt without payload is an engine failure");
+    }
 }
@@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{
    ChatMessageDelta,
 };
 use crate::routes::openai::utils::logprobs::{
-    decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
+    decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
 };
 use crate::routes::openai::utils::types::{
    ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
@@ -181,14 +181,11 @@ async fn collect_chat_completion(
        None
    };
    let prompt_logprobs = if include_prompt_logprobs {
-        Some(decoded_prompt_logprobs_to_maps(
-            prompt_logprobs.as_ref().ok_or_else(|| {
-                server_error!(
-                    "chat response requested prompt_logprobs but generation returned none"
-                )
-            })?,
+        Some(prompt_logprobs_to_maps(
+            prompt_logprobs.as_ref(),
+            &prompt_token_ids,
            return_tokens_as_token_ids,
-        ))
+        )?)
    } else {
        None
    };
@@ -5,7 +5,6 @@ mod convert;
 mod types;
 mod validate;

-use std::collections::HashMap;
 use std::convert::Infallible;
 use std::result::Result;
 use std::sync::Arc;
@@ -29,8 +28,8 @@ use vllm_text::{

 use self::convert::{ResponseOptions, prepare_completion_request};
 use super::utils::logprobs::{
-    collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
-    decoded_prompt_logprobs_to_openai, text_len,
+    collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
+    prompt_logprobs_to_maps, text_len,
 };
 use super::utils::types::Usage;
 use crate::config::ApiServerOptions;
@@ -505,27 +504,6 @@ fn prompt_only_logprobs_to_openai(
    ))
 }

-fn prompt_logprobs_to_maps(
-    prompt_logprobs: Option<&DecodedPromptLogprobs>,
-    prompt_token_ids: &[u32],
-    return_tokens_as_token_ids: bool,
-) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
-    if let Some(prompt_logprobs) = prompt_logprobs {
-        return Ok(decoded_prompt_logprobs_to_maps(
-            prompt_logprobs,
-            return_tokens_as_token_ids,
-        ));
-    }
-
-    if let [_token_id] = prompt_token_ids {
-        return Ok(vec![None]);
-    }
-
-    Err(server_error!(
-        "completion response requested prompt_logprobs but generation returned none"
-    ))
-}
-
 fn usage_chunk(
    request_id: &str,
    response_model: &str,
@@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai(
    })
 }

-/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
-/// shape.
-pub fn decoded_prompt_logprobs_to_maps(
-    prompt_logprobs: &DecodedPromptLogprobs,
+/// Map decoded prompt logprobs into vLLM-style per-position maps, treating a
+/// missing single-token payload as `[None]`.
+pub fn prompt_logprobs_to_maps(
+    prompt_logprobs: Option<&DecodedPromptLogprobs>,
+    prompt_token_ids: &[u32],
    return_tokens_as_token_ids: bool,
-) -> Vec<Option<HashMap<String, f32>>> {
-    std::iter::once(None)
-        .chain(prompt_logprobs.scored_positions.iter().map(|position| {
-            Some(position_top_logprobs_map(
-                position,
-                return_tokens_as_token_ids,
-            ))
-        }))
-        .collect()
+) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
+    if let Some(prompt_logprobs) = prompt_logprobs {
+        return Ok(std::iter::once(None)
+            .chain(prompt_logprobs.scored_positions.iter().map(|position| {
+                Some(position_top_logprobs_map(
+                    position,
+                    return_tokens_as_token_ids,
+                ))
+            }))
+            .collect());
+    }
+
+    if let [_token_id] = prompt_token_ids {
+        return Ok(vec![None]);
+    }
+
+    Err(server_error!(
+        "prompt_logprobs were requested but generation returned none"
+    ))
 }

 /// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
@@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
 mod tests {
    use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};

-    use super::decoded_logprobs_to_openai_chat;
+    use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps};
+
+    #[test]
+    fn prompt_logprobs_maps_reject_missing_multi_token_payload() {
+        prompt_logprobs_to_maps(None, &[9707, 11], false)
+            .expect_err("multi-token prompt without payload is an engine failure");
+    }

    fn sample_logprobs() -> DecodedLogprobs {
        DecodedLogprobs {
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+from types import ModuleType, SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vllm.platforms.interface import DeviceCapability
+
+
+@pytest.fixture
+def cuda_platform_base(monkeypatch: pytest.MonkeyPatch) -> Any:
+    stable_libtorch_module = ModuleType("vllm._C_stable_libtorch")
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm._C_stable_libtorch",
+        stable_libtorch_module,
+    )
+    from vllm.platforms.cuda import CudaPlatformBase
+
+    return CudaPlatformBase
+
+
+def test_compiled_arch_covers_device(cuda_platform_base: Any) -> None:
+    assert cuda_platform_base._compiled_arch_covers_device(
+        "12.1a", DeviceCapability(12, 1)
+    )
+    assert not cuda_platform_base._compiled_arch_covers_device(
+        "12.0a", DeviceCapability(12, 1)
+    )
+
+    assert cuda_platform_base._compiled_arch_covers_device(
+        "12.0f", DeviceCapability(12, 1)
+    )
+    assert not cuda_platform_base._compiled_arch_covers_device(
+        "10.0f", DeviceCapability(12, 1)
+    )
+
+
+def test_warn_if_device_arch_not_compiled(
+    monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
+) -> None:
+    def device_count(cls: type[Any]) -> int:
+        return 2
+
+    def get_device_capability(
+        cls: type[Any], device_id: int = 0
+    ) -> DeviceCapability | None:
+        capabilities = {
+            0: DeviceCapability(12, 1),
+            1: DeviceCapability(10, 3),
+        }
+        return capabilities[device_id]
+
+    def get_device_name(cls: type[Any], device_id: int = 0) -> str:
+        return f"GPU {device_id}"
+
+    warnings: list[tuple[str, str, str]] = []
+
+    def warning_once(message: str, compiled_archs: str, devices: str) -> None:
+        warnings.append((message, compiled_archs, devices))
+
+    monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
+    monkeypatch.setattr(
+        cuda_platform_base,
+        "get_device_capability",
+        classmethod(get_device_capability),
+    )
+    monkeypatch.setattr(
+        cuda_platform_base, "get_device_name", classmethod(get_device_name)
+    )
+    monkeypatch.setattr(
+        "vllm.platforms.cuda.torch.ops",
+        SimpleNamespace(
+            _C=SimpleNamespace(get_compiled_cuda_archs=lambda: "12.0f,10.0a")
+        ),
+    )
+    monkeypatch.setattr(
+        "vllm.platforms.cuda.logger",
+        SimpleNamespace(warning_once=warning_once),
+    )
+
+    cuda_platform_base._warn_if_device_arch_not_compiled()
+
+    assert len(warnings) == 1
+    assert warnings[0][1] == "12.0f, 10.0a"
+    assert "1: GPU 1 (compute capability 10.3)" in warnings[0][2]
+
+
+def test_warn_if_device_arch_not_compiled_no_warning(
+    monkeypatch: pytest.MonkeyPatch, cuda_platform_base: Any
+) -> None:
+    def device_count(cls: type[Any]) -> int:
+        return 1
+
+    def get_device_capability(
+        cls: type[Any], device_id: int = 0
+    ) -> DeviceCapability | None:
+        return DeviceCapability(10, 3)
+
+    def get_device_name(cls: type[Any], device_id: int = 0) -> str:
+        raise AssertionError("get_device_name should not be called for covered devices")
+
+    warnings: list[tuple[str, str, str]] = []
+
+    def warning_once(message: str, compiled_archs: str, devices: str) -> None:
+        warnings.append((message, compiled_archs, devices))
+
+    monkeypatch.setattr(cuda_platform_base, "device_count", classmethod(device_count))
+    monkeypatch.setattr(
+        cuda_platform_base,
+        "get_device_capability",
+        classmethod(get_device_capability),
+    )
+    monkeypatch.setattr(
+        cuda_platform_base, "get_device_name", classmethod(get_device_name)
+    )
+    monkeypatch.setattr(
+        "vllm.platforms.cuda.torch.ops",
+        SimpleNamespace(_C=SimpleNamespace(get_compiled_cuda_archs=lambda: "10.0f")),
+    )
+    monkeypatch.setattr(
+        "vllm.platforms.cuda.logger",
+        SimpleNamespace(warning_once=warning_once),
+    )
+
+    cuda_platform_base._warn_if_device_arch_not_compiled()
+
+    assert warnings == []
@@ -60,6 +60,8 @@ def _make_quick_allreduce(
    qar.use_fp16_kernels = use_fp16_kernels
    qar.qr_quant_level = QuickReduceRegime[quant_level]
    qar.qr_max_size = qr_max_size
+    qar.qr_min_size = None
+    qar.qr_quantization_min_size = None
    return qar


@@ -511,13 +513,21 @@ def test_quick_reduce_regime_values():
    assert QuickReduceRegime.INT8.value == 1
    assert QuickReduceRegime.INT6.value == 2
    assert QuickReduceRegime.INT4.value == 3
-    assert QuickReduceRegime.NONE.value == 4
+    assert QuickReduceRegime.INT3.value == 4
+    assert QuickReduceRegime.NONE.value == 5


 def test_quick_reduce_regime_names():
    from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime

-    assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
+    assert set(QuickReduceRegime.__members__) == {
+        "FP",
+        "INT8",
+        "INT6",
+        "INT4",
+        "INT3",
+        "NONE",
+    }


@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
@@ -693,7 +703,7 @@ def test_quick_allreduce_min_size_table():
    for dtype in [torch.float16, torch.bfloat16]:
        for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
            min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
-            assert len(min_sizes) == 4
+            assert len(min_sizes) == 5
            assert all(size > 0 for size in min_sizes)


@@ -515,3 +515,15 @@ def test_structured_outputs_structural_tag_invalid(structural_tag):
            messages=[{"role": "user", "content": "hello"}],
            structured_outputs={"structural_tag": structural_tag},
        )
+
+
+@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"])
+def test_non_numeric_logprobs_rejected(field_name):
+    """A non-numeric logprobs value must be a clean 400 validation error, not a
+    TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
+    with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
+        ChatCompletionRequest(
+            model=MODEL_NAME,
+            messages=[{"role": "user", "content": "hello"}],
+            **{field_name: "2"},
+        )
@@ -610,3 +610,16 @@ class TestCompletionPromptListLimit:
            max_tokens=1,
        )
        assert len(request.prompt_embeds) == 5
+
+
+@pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"])
+def test_non_numeric_logprobs_rejected(field_name):
+    """A non-numeric logprobs value must be a clean 400 validation error, not a
+    TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
+    with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
+        CompletionRequest(
+            model=MODEL_NAME,
+            prompt="Test prompt",
+            max_tokens=10,
+            **{field_name: "2"},
+        )
@@ -425,7 +425,7 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No
    match the single-call result.
    """
    from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
-        causal_conv1d_torch,
+        causal_conv1d_fn_cpu as causal_conv1d_torch,
    )

    x, weight, bias = _conv_inputs(total_tokens)
@@ -18,8 +18,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 DEVICE = current_platform.device_type

 pytestmark = pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
-    reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
+    not (
+        current_platform.is_cuda_alike()
+        or current_platform.is_xpu()
+        or current_platform.is_cpu()
+    ),
+    reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU",
 )


@@ -284,7 +288,8 @@ def test_causal_conv1d_varlen(
    batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
 ):
    device = DEVICE
-    torch.accelerator.empty_cache()
+    if not current_platform.is_cpu():
+        torch.accelerator.empty_cache()
    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
    if itype == torch.bfloat16:
        rtol, atol = 1e-2, 5e-2
@@ -20,8 +20,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 DEVICE = current_platform.device_type

 pytestmark = pytest.mark.skipif(
-    not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
-    reason="mamba_ssm kernels require CUDA-alike or XPU",
+    not (
+        current_platform.is_cuda_alike()
+        or current_platform.is_xpu()
+        or current_platform.is_cpu()
+    ),
+    reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU",
 )

 # selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
@@ -342,6 +346,13 @@ def test_selective_scan(
@pytest.mark.parametrize("has_z", [False, True])
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update(dim, dstate, has_z, itype):
    device = DEVICE
    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -436,6 +447,13 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
@pytest.mark.parametrize("max_seq_len", [1, 2, 4])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
    device = DEVICE
    rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -697,6 +715,13 @@ def test_selective_scan_varlen(
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
 # tests correctness in case subset of the sequences are padded
@pytest.mark.parametrize("with_padding", [True, False])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_with_batch_indices(
    with_padding, dim, dstate, has_z, itype
 ):
@@ -789,6 +814,13 @@ def test_selective_state_update_with_batch_indices(
@pytest.mark.parametrize("ngroups", [1, 4])
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_with_heads_with_batch_indices(
    dim, dstate, ngroups, has_z, tie_hdim, itype
 ):
@@ -862,6 +894,13 @@ def test_selective_state_update_with_heads_with_batch_indices(
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.parametrize("max_seq_len", [2, 4])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_with_num_accepted_tokens(
    dim, dstate, has_z, itype, max_seq_len
 ):
@@ -988,6 +1027,13 @@ def test_selective_state_update_with_num_accepted_tokens(
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.parametrize("max_seq_len", [2, 4])
+@pytest.mark.skipif(
+    current_platform.is_cpu(),
+    reason=(
+        "CPU kernel for selective_state_update only supports "
+        "Mamba 2 (scalar A/dt), not Mamba 1."
+    ),
+)
 def test_selective_state_update_varlen_with_num_accepted(
    dim, dstate, has_z, itype, max_seq_len
 ):
@@ -0,0 +1,880 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end tests for routed-expert capture on the monolithic MoE path.
+
+These tests exercise the wiring that lets ``RoutedExpertsCapturer`` see the
+expert IDs picked by FlashInfer's fused router-and-experts kernels (the
+"monolithic" path). When ``set_capture_fn`` is installed on
+a ``FusedMoEExpertsMonolithic`` subclass that supports it, the kernel call
+should:
+
+  * allocate an int16 ``(num_tokens, top_k)`` buffer,
+  * pass it to FlashInfer as ``routing_replay_out``,
+  * invoke the callback after the kernel returns.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEParallelConfig,
+    FusedMoEQuantConfig,
+    RoutingMethodType,
+    fp8_w8a8_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
+    TrtLlmBf16ExpertsMonolithic,
+)
+from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
+    TrtLlmFp8ExpertsMonolithic,
+)
+from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
+    TrtLlmNvFp4ExpertsMonolithic,
+)
+from vllm.platforms import current_platform
+
+try:
+    from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
+except ImportError:
+    pytest.skip("flashinfer not available", allow_module_level=True)
+
+if not has_flashinfer_trtllm_fused_moe() or not current_platform.is_cuda():
+    pytest.skip(
+        "Requires FlashInfer TRT-LLM fused MoE on CUDA",
+        allow_module_level=True,
+    )
+
+if not current_platform.has_device_capability(100):
+    pytest.skip(
+        "TRT-LLM fused MoE kernels require SM100+",
+        allow_module_level=True,
+    )
+
+
+def _shuffle_bf16_weights_block_major_k(
+    w: torch.Tensor, epilogue_tile_m: int = 64, block_k: int = 128
+) -> torch.Tensor:
+    """Reshape ``w`` (E, M, K) into the ``BlockMajorK`` layout expected by
+    ``trtllm_bf16_moe``: ``(E, K/block_k, M, block_k)`` after a per-expert
+    row shuffle.
+    """
+    from flashinfer import shuffle_matrix_a
+    from flashinfer.fused_moe import convert_to_block_layout
+
+    num_experts = w.shape[0]
+    shuffled = []
+    for i in range(num_experts):
+        t = shuffle_matrix_a(w[i].view(torch.uint8), epilogue_tile_m)
+        shuffled.append(convert_to_block_layout(t, block_k))
+    return torch.stack(shuffled).view(torch.bfloat16)
+
+
+def _make_bf16_monolithic_experts(
+    num_experts: int,
+    top_k: int,
+    hidden_size: int,
+    intermediate_size: int,
+    routing_method: RoutingMethodType,
+    device: torch.device,
+) -> tuple[TrtLlmBf16ExpertsMonolithic, torch.Tensor, torch.Tensor]:
+    """Construct the monolithic BF16 experts plus the BlockMajorK weights
+    expected by ``trtllm_bf16_moe``.
+    """
+    parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
+    moe_config = FusedMoEConfig(
+        num_experts=num_experts,
+        experts_per_token=top_k,
+        hidden_dim=hidden_size,
+        intermediate_size=intermediate_size,
+        num_local_experts=num_experts,
+        num_logical_experts=num_experts,
+        moe_parallel_config=parallel_cfg,
+        in_dtype=torch.bfloat16,
+        activation=MoEActivation.SILU,
+        device=device,
+        routing_method=routing_method,
+        max_num_tokens=max(8, 1),
+    )
+    quant_config = FusedMoEQuantConfig.make(
+        quant_dtype=None,
+        per_act_token_quant=False,
+        per_out_ch_quant=False,
+        block_shape=None,
+    )
+
+    experts = TrtLlmBf16ExpertsMonolithic(
+        moe_config=moe_config, quant_config=quant_config
+    )
+
+    gemm1 = (
+        torch.randn(
+            num_experts,
+            2 * intermediate_size,
+            hidden_size,
+            device=device,
+            dtype=torch.bfloat16,
+        )
+        * 0.1
+    )
+    gemm2 = (
+        torch.randn(
+            num_experts,
+            hidden_size,
+            intermediate_size,
+            device=device,
+            dtype=torch.bfloat16,
+        )
+        * 0.1
+    )
+    w13 = _shuffle_bf16_weights_block_major_k(gemm1)
+    w2 = _shuffle_bf16_weights_block_major_k(gemm2)
+    return experts, w13, w2
+
+
+def _run_bf16_monolithic(
+    experts: TrtLlmBf16ExpertsMonolithic,
+    hidden_states: torch.Tensor,
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    router_logits: torch.Tensor,
+    num_experts: int,
+    *,
+    n_group: int | None = None,
+    topk_group: int | None = None,
+    routed_scaling_factor: float | None = None,
+    e_score_correction_bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    return experts.apply(
+        hidden_states=hidden_states,
+        w1=w13,
+        w2=w2,
+        router_logits=router_logits,
+        activation=MoEActivation.SILU,
+        global_num_experts=num_experts,
+        expert_map=None,
+        a1q_scale=None,
+        apply_router_weight_on_input=False,
+        num_expert_group=n_group,
+        topk_group=topk_group,
+        e_score_correction_bias=e_score_correction_bias,
+        routed_scaling_factor=routed_scaling_factor,
+    )
+
+
+_DSV3_NUM_EXPERTS = 32
+_DSV3_N_GROUP = 4
+_DSV3_TOPK_GROUP = 2
+
+
+def _make_dsv3_routing_bias(num_experts: int, device: torch.device) -> torch.Tensor:
+    return torch.randn(num_experts, device=device, dtype=torch.bfloat16)
+
+
+@pytest.mark.parametrize("num_tokens", [2, 7, 16])
+@pytest.mark.parametrize("top_k", [2, 4])
+def test_trtllm_bf16_monolithic_routing_replay_records_valid_experts(
+    num_tokens: int,
+    top_k: int,
+) -> None:
+    """The capture callback should receive the int16 routed-expert IDs the
+    kernel actually used, the values should be valid expert indices, and
+    each token should pick ``top_k`` distinct experts."""
+    if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
+        pytest.skip(
+            f"DSV3 requires top_k <= n_group * topk_group "
+            f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
+        )
+    torch.manual_seed(0)
+    device = torch.device("cuda:0")
+
+    num_experts = _DSV3_NUM_EXPERTS
+    hidden_size = 1024
+    intermediate_size = 1024
+
+    experts, w13, w2 = _make_bf16_monolithic_experts(
+        num_experts=num_experts,
+        top_k=top_k,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        routing_method=RoutingMethodType.DeepSeekV3,
+        device=device,
+    )
+
+    captured: list[torch.Tensor] = []
+
+    def capture_fn(replay_out: torch.Tensor) -> None:
+        captured.append(replay_out.clone())
+
+    assert experts.supports_routing_replay_capture()
+    experts.set_capture_fn(capture_fn)
+
+    hidden_states = (
+        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
+    )
+    router_logits = torch.rand(
+        num_tokens, num_experts, device=device, dtype=torch.float32
+    )
+    routing_bias = _make_dsv3_routing_bias(num_experts, device)
+
+    _ = _run_bf16_monolithic(
+        experts,
+        hidden_states=hidden_states,
+        w13=w13,
+        w2=w2,
+        router_logits=router_logits,
+        num_experts=num_experts,
+        n_group=_DSV3_N_GROUP,
+        topk_group=_DSV3_TOPK_GROUP,
+        routed_scaling_factor=1.0,
+        e_score_correction_bias=routing_bias,
+    )
+
+    assert len(captured) == 1
+    replay = captured[0]
+    assert replay.dtype == torch.int16
+    assert replay.shape == (num_tokens, top_k)
+    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
+    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
+
+    for t in range(num_tokens):
+        unique = replay[t].unique()
+        assert unique.numel() == top_k, (
+            f"token {t}: expected {top_k} distinct experts, "
+            f"got {unique.numel()} ({replay[t].tolist()})"
+        )
+
+
+@pytest.mark.parametrize("num_tokens", [2, 7, 16])
+@pytest.mark.parametrize(
+    "routing_method",
+    [
+        RoutingMethodType.Renormalize,
+        RoutingMethodType.RenormalizeNaive,
+    ],
+)
+def test_trtllm_bf16_monolithic_routing_replay_non_dsv3(
+    num_tokens: int,
+    routing_method: RoutingMethodType,
+) -> None:
+    """Routing replay works for non-DeepSeekV3 routing methods too.
+    FlashInfer's ``routing_replay_out`` is routing-method-agnostic."""
+    torch.manual_seed(0)
+    device = torch.device("cuda:0")
+
+    num_experts = 8
+    top_k = 2
+    hidden_size = 1024
+    intermediate_size = 1024
+
+    experts, w13, w2 = _make_bf16_monolithic_experts(
+        num_experts=num_experts,
+        top_k=top_k,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        routing_method=routing_method,
+        device=device,
+    )
+
+    captured: list[torch.Tensor] = []
+    experts.set_capture_fn(lambda r: captured.append(r.clone()))
+
+    hidden_states = (
+        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
+    )
+    router_logits = torch.rand(
+        num_tokens, num_experts, device=device, dtype=torch.float32
+    )
+
+    _ = _run_bf16_monolithic(
+        experts,
+        hidden_states=hidden_states,
+        w13=w13,
+        w2=w2,
+        router_logits=router_logits,
+        num_experts=num_experts,
+    )
+
+    assert len(captured) == 1
+    replay = captured[0]
+    assert replay.dtype == torch.int16
+    assert replay.shape == (num_tokens, top_k)
+    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
+    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
+    for t in range(num_tokens):
+        unique = replay[t].unique()
+        assert unique.numel() == top_k, (
+            f"token {t}: expected {top_k} distinct experts, "
+            f"got {unique.numel()} ({replay[t].tolist()})"
+        )
+
+
+def test_trtllm_bf16_monolithic_capture_disabled_skips_buffer_alloc() -> None:
+    """With no callback installed the kernel should not see a
+    ``routing_replay_out`` tensor — verify the helper short-circuits."""
+    torch.manual_seed(0)
+    device = torch.device("cuda:0")
+    experts, _, _ = _make_bf16_monolithic_experts(
+        num_experts=_DSV3_NUM_EXPERTS,
+        top_k=2,
+        hidden_size=1024,
+        intermediate_size=1024,
+        routing_method=RoutingMethodType.DeepSeekV3,
+        device=device,
+    )
+    # No callback installed.
+    buf = experts._maybe_make_routing_replay_buffer(num_tokens=4, device=device)
+    assert buf is None
+
+    # Dispatch is also a no-op.
+    experts._maybe_dispatch_routing_replay(buf, num_tokens=4)
+
+
+def test_trtllm_bf16_monolithic_supports_capture_for_all_routing() -> None:
+    """FlashInfer's ``routing_replay_out`` is supported by all routing
+    methods, so ``supports_routing_replay_capture`` should be True
+    regardless of routing method."""
+    device = torch.device("cuda:0")
+    for routing_method in (
+        RoutingMethodType.DeepSeekV3,
+        RoutingMethodType.Renormalize,
+        RoutingMethodType.RenormalizeNaive,
+    ):
+        experts, _, _ = _make_bf16_monolithic_experts(
+            num_experts=_DSV3_NUM_EXPERTS,
+            top_k=2,
+            hidden_size=1024,
+            intermediate_size=1024,
+            routing_method=routing_method,
+            device=device,
+        )
+        assert experts.supports_routing_replay_capture() is True, (
+            f"{routing_method!r} should support routing replay capture"
+        )
+
+
+def test_trtllm_bf16_monolithic_capture_buffer_shape_and_dtype() -> None:
+    """When capture is installed, the allocated buffer is int16 and shaped
+    ``(num_tokens, experts_per_token)``."""
+    device = torch.device("cuda:0")
+    experts, _, _ = _make_bf16_monolithic_experts(
+        num_experts=_DSV3_NUM_EXPERTS,
+        top_k=4,
+        hidden_size=1024,
+        intermediate_size=1024,
+        routing_method=RoutingMethodType.DeepSeekV3,
+        device=device,
+    )
+    experts.set_capture_fn(lambda r: None)
+    buf = experts._maybe_make_routing_replay_buffer(num_tokens=11, device=device)
+    assert buf is not None
+    assert buf.dtype == torch.int16
+    assert buf.shape[0] >= 11
+    assert buf.shape[1] == 4
+    assert buf.device.type == "cuda"
+
+
+def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
+    """End-to-end: bind ``RoutedExpertsCapturer.capture`` as the callback
+    on the monolithic experts and verify the captured rows land in the
+    capturer's device buffer at the correct layer slot.
+
+    Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
+    for the monolithic path: a single closure is installed on the monolithic
+    ``fused_experts`` (in addition to ``router.set_capture_fn`` on the
+    non-monolithic path) and the capturer routes per-layer based on the
+    closed-over ``layer_id``.
+    """
+    from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+        RoutedExpertsCapturer,
+    )
+
+    torch.manual_seed(7)
+    device = torch.device("cuda:0")
+    num_tokens = 4
+    top_k = 2
+    num_experts = _DSV3_NUM_EXPERTS
+    hidden_size = 1024
+    intermediate_size = 1024
+
+    experts, w13, w2 = _make_bf16_monolithic_experts(
+        num_experts=num_experts,
+        top_k=top_k,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        routing_method=RoutingMethodType.DeepSeekV3,
+        device=device,
+    )
+
+    num_layers = 3
+    layer_id = 1
+    capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
+    capturer.dp_rank = 0
+    capturer.tp_size = 1
+    capturer.device_buffer = torch.full(
+        (num_tokens + 4, num_layers, top_k),
+        -1,
+        dtype=torch.int32,
+        device=device,
+    )
+
+    def capture_fn(replay_out: torch.Tensor) -> None:
+        capturer.capture(layer_id, replay_out)
+
+    experts.set_capture_fn(capture_fn)
+
+    hidden_states = (
+        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
+    )
+    router_logits = torch.rand(
+        num_tokens, num_experts, device=device, dtype=torch.float32
+    )
+    routing_bias = _make_dsv3_routing_bias(num_experts, device)
+
+    # Patch get_forward_context to return a dp_metadata=None context so the
+    # capturer takes the single-DP branch.
+    import vllm.model_executor.layers.fused_moe.routed_experts_capturer as rec
+
+    with patch.object(
+        rec,
+        "get_forward_context",
+        return_value=SimpleNamespace(dp_metadata=None),
+    ):
+        _ = _run_bf16_monolithic(
+            experts,
+            hidden_states=hidden_states,
+            w13=w13,
+            w2=w2,
+            router_logits=router_logits,
+            num_experts=num_experts,
+            n_group=_DSV3_N_GROUP,
+            topk_group=_DSV3_TOPK_GROUP,
+            routed_scaling_factor=1.0,
+            e_score_correction_bias=routing_bias,
+        )
+
+    captured = capturer.device_buffer[:num_tokens, layer_id, :].cpu()
+    # Valid expert IDs at this layer.
+    assert (captured >= 0).all()
+    assert (captured < num_experts).all()
+    for t in range(num_tokens):
+        unique = captured[t].unique()
+        assert unique.numel() == top_k, (
+            f"token {t}: expected {top_k} distinct experts at layer "
+            f"{layer_id}, got {unique.numel()}"
+        )
+
+    # Other layers / trailing token rows untouched.
+    for other_layer in range(num_layers):
+        if other_layer == layer_id:
+            continue
+        assert (capturer.device_buffer[:, other_layer, :].cpu() == -1).all(), (
+            f"layer {other_layer} should be untouched, got writes"
+        )
+    assert (capturer.device_buffer[num_tokens:, layer_id, :].cpu() == -1).all(), (
+        "tail rows beyond num_tokens should remain sentinel"
+    )
+
+
+# ----------------------------------------------------------------------------
+# FP8 block-scale (DeepSeekFp8) — vLLM's ``TrtLlmFp8ExpertsMonolithic``
+# ----------------------------------------------------------------------------
+
+
+def _make_fp8_block_scale_monolithic_experts(
+    num_experts: int,
+    top_k: int,
+    hidden_size: int,
+    intermediate_size: int,
+    device: torch.device,
+) -> tuple[TrtLlmFp8ExpertsMonolithic, torch.Tensor, torch.Tensor]:
+    """Set up ``TrtLlmFp8ExpertsMonolithic`` for the DeepSeekFp8 block-scale
+    code path with DSV3 routing.
+
+    Weights are shuffled into the BlockMajorK layout the kernel expects
+    (same helper the vLLM weight loader uses for DeepSeek-FP8 models).
+    """
+    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
+        _shuffle_deepseek_fp8_moe_weights,
+    )
+
+    block_k = 128
+    parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
+    moe_config = FusedMoEConfig(
+        num_experts=num_experts,
+        experts_per_token=top_k,
+        hidden_dim=hidden_size,
+        intermediate_size=intermediate_size,
+        num_local_experts=num_experts,
+        num_logical_experts=num_experts,
+        moe_parallel_config=parallel_cfg,
+        in_dtype=torch.bfloat16,
+        activation=MoEActivation.SILU,
+        device=device,
+        routing_method=RoutingMethodType.DeepSeekV3,
+        max_num_tokens=max(8, 1),
+    )
+
+    # Random fp8 weights + ones-block scales (the kernel decoder only cares
+    # that the per-block scales are present and finite for routing/replay).
+    gemm1 = torch.randn(
+        num_experts, 2 * intermediate_size, hidden_size, device=device
+    ).to(torch.float8_e4m3fn)
+    gemm2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device).to(
+        torch.float8_e4m3fn
+    )
+    w13_shuffled, w2_shuffled = _shuffle_deepseek_fp8_moe_weights(gemm1, gemm2)
+
+    w1_scale = torch.ones(
+        num_experts,
+        2 * intermediate_size // block_k,
+        hidden_size // block_k,
+        device=device,
+        dtype=torch.float32,
+    )
+    w2_scale = torch.ones(
+        num_experts,
+        hidden_size // block_k,
+        intermediate_size // block_k,
+        device=device,
+        dtype=torch.float32,
+    )
+    quant_config = fp8_w8a8_moe_quant_config(
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        block_shape=[block_k, block_k],
+        per_act_token_quant=False,
+    )
+
+    experts = TrtLlmFp8ExpertsMonolithic(
+        moe_config=moe_config, quant_config=quant_config
+    )
+    return experts, w13_shuffled, w2_shuffled
+
+
+def _run_fp8_block_scale_monolithic(
+    experts: TrtLlmFp8ExpertsMonolithic,
+    hidden_states_fp8: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    router_logits: torch.Tensor,
+    num_experts: int,
+    routing_bias: torch.Tensor,
+) -> torch.Tensor:
+    return experts.apply(
+        hidden_states=hidden_states_fp8,
+        w1=w13,
+        w2=w2,
+        router_logits=router_logits,
+        activation=MoEActivation.SILU,
+        global_num_experts=num_experts,
+        expert_map=None,
+        # The block-scale apply path reads ``a1q_scale`` and transposes it
+        # to ``(hidden_size/128, num_tokens)`` for the kernel call.
+        a1q_scale=hidden_states_scale,
+        apply_router_weight_on_input=False,
+        num_expert_group=_DSV3_N_GROUP,
+        topk_group=_DSV3_TOPK_GROUP,
+        e_score_correction_bias=routing_bias,
+        routed_scaling_factor=1.0,
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [2, 7, 16])
+@pytest.mark.parametrize("top_k", [2, 4])
+def test_trtllm_fp8_block_scale_monolithic_routing_replay_records_valid_experts(
+    num_tokens: int,
+    top_k: int,
+) -> None:
+    """End-to-end: ``TrtLlmFp8ExpertsMonolithic`` (DeepSeekFp8 block-scale
+    path, DSV3 routing) captures valid expert IDs."""
+    if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
+        pytest.skip(
+            f"DSV3 requires top_k <= n_group * topk_group "
+            f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
+        )
+    torch.manual_seed(0)
+    device = torch.device("cuda:0")
+
+    num_experts = _DSV3_NUM_EXPERTS
+    hidden_size = 1024
+    intermediate_size = 1024
+
+    experts, w13, w2 = _make_fp8_block_scale_monolithic_experts(
+        num_experts=num_experts,
+        top_k=top_k,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        device=device,
+    )
+    assert experts.supports_routing_replay_capture()
+
+    captured: list[torch.Tensor] = []
+    experts.set_capture_fn(lambda r: captured.append(r.clone()))
+
+    # Per-token / per-block hidden scales (ones is fine for the routing
+    # path; the GEMM output isn't being asserted on).
+    hidden_states = (
+        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
+    ).to(torch.float8_e4m3fn)
+    hidden_states_scale = torch.ones(
+        num_tokens, hidden_size // 128, device=device, dtype=torch.float32
+    )
+    router_logits = torch.rand(
+        num_tokens, num_experts, device=device, dtype=torch.float32
+    )
+    routing_bias = _make_dsv3_routing_bias(num_experts, device)
+
+    _ = _run_fp8_block_scale_monolithic(
+        experts,
+        hidden_states_fp8=hidden_states,
+        hidden_states_scale=hidden_states_scale,
+        w13=w13,
+        w2=w2,
+        router_logits=router_logits,
+        num_experts=num_experts,
+        routing_bias=routing_bias,
+    )
+
+    assert len(captured) == 1
+    replay = captured[0]
+    assert replay.dtype == torch.int16
+    assert replay.shape == (num_tokens, top_k)
+    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
+    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
+    for t in range(num_tokens):
+        unique = replay[t].unique()
+        assert unique.numel() == top_k, (
+            f"token {t}: expected {top_k} distinct experts, "
+            f"got {unique.numel()} ({replay[t].tolist()})"
+        )
+
+
+# ----------------------------------------------------------------------------
+# NVFP4 — vLLM's ``TrtLlmNvFp4ExpertsMonolithic``
+# ----------------------------------------------------------------------------
+
+
+def _make_nvfp4_monolithic_experts(
+    num_experts: int,
+    top_k: int,
+    hidden_size: int,
+    intermediate_size: int,
+    device: torch.device,
+) -> tuple[
+    TrtLlmNvFp4ExpertsMonolithic,
+    torch.Tensor,  # w13 (packed nvfp4 uint8)
+    torch.Tensor,  # w13 block-scale (fp8)
+    torch.Tensor,  # w2  (packed nvfp4 uint8)
+    torch.Tensor,  # w2  block-scale (fp8)
+    torch.Tensor,  # input global scale (per-tensor float32)
+]:
+    """Set up ``TrtLlmNvFp4ExpertsMonolithic`` with NVFP4-quantized weights
+    and DSV3 routing.
+
+    NVFP4 = per-block-of-16 fp4 with an fp8 scale, plus a per-tensor
+    "global" scale. We follow the layout in ``test_ocp_mx_moe.py`` /
+    ``flashinfer/tests/moe/test_trtllm_gen_routed_fused_moe.py``:
+      * weights: uint8 (packed fp4) ``(E, M, K//2)``
+      * weight scales: fp8 ``(E, M, K//16)``
+      * hidden states: uint8 (packed fp4) ``(N, K//2)``
+      * hidden state scales: fp8 ``(N, K//16)``
+    """
+    from flashinfer import fp4_quantize
+
+    block_size = 16
+    parallel_cfg = FusedMoEParallelConfig.make_no_parallel()
+    moe_config = FusedMoEConfig(
+        num_experts=num_experts,
+        experts_per_token=top_k,
+        hidden_dim=hidden_size,
+        intermediate_size=intermediate_size,
+        num_local_experts=num_experts,
+        num_logical_experts=num_experts,
+        moe_parallel_config=parallel_cfg,
+        in_dtype=torch.bfloat16,
+        activation=MoEActivation.SILU,
+        device=device,
+        routing_method=RoutingMethodType.DeepSeekV3,
+        max_num_tokens=max(8, 1),
+    )
+
+    gemm1 = torch.randn(
+        num_experts,
+        2 * intermediate_size,
+        hidden_size,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    gemm2 = torch.randn(
+        num_experts,
+        hidden_size,
+        intermediate_size,
+        device=device,
+        dtype=torch.bfloat16,
+    )
+    # Per-tensor weight scaling factor (used to build ``g1_alphas`` /
+    # ``g2_alphas`` below).
+    w_global_scale = torch.tensor(1.0, device=device)
+    # Per-tensor input scaling factor.
+    a_global_scale = torch.tensor(1.0, device=device)
+
+    w13_q, w13_scale = fp4_quantize(
+        gemm1,
+        w_global_scale,
+        block_size,
+        sf_use_ue8m0=False,
+        is_sf_swizzled_layout=False,
+    )
+    w13_scale = w13_scale.view(torch.float8_e4m3fn).reshape(
+        num_experts, 2 * intermediate_size, hidden_size // block_size
+    )
+    w2_q, w2_scale = fp4_quantize(
+        gemm2,
+        w_global_scale,
+        block_size,
+        sf_use_ue8m0=False,
+        is_sf_swizzled_layout=False,
+    )
+    w2_scale = w2_scale.view(torch.float8_e4m3fn).reshape(
+        num_experts, hidden_size, intermediate_size // block_size
+    )
+
+    # NVFP4 dq scale chain: g1_alphas = w1_scale_2 * a1_scale_2,
+    # g2_alphas = w2_scale_2 * a2_scale_2. The kernel multiplies by these.
+    g_alphas = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
+    a2_gscale = torch.full((num_experts,), 1.0, device=device, dtype=torch.float32)
+
+    quant_config = FusedMoEQuantConfig.make(
+        quant_dtype="nvfp4",
+        per_act_token_quant=False,
+        per_out_ch_quant=False,
+        block_shape=None,
+        w1_scale=w13_scale,
+        w2_scale=w2_scale,
+        g1_alphas=g_alphas,
+        g2_alphas=g_alphas,
+        a1_gscale=a_global_scale,
+        a2_gscale=a2_gscale,
+    )
+
+    experts = TrtLlmNvFp4ExpertsMonolithic(
+        moe_config=moe_config, quant_config=quant_config
+    )
+    return experts, w13_q, w13_scale, w2_q, w2_scale, a_global_scale
+
+
+def _run_nvfp4_monolithic(
+    experts: TrtLlmNvFp4ExpertsMonolithic,
+    hidden_states_q: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    router_logits: torch.Tensor,
+    num_experts: int,
+    routing_bias: torch.Tensor,
+) -> torch.Tensor:
+    """The monolithic NVFP4 apply expects packed fp4 hidden states + the
+    matching fp8 per-block scale stored in the ``a1q_scale`` slot."""
+    # Stash the weight tensors on the experts in the locations the apply()
+    # implementation reads from (it pulls them from quant_config / scales
+    # already; w1/w2 come in as args).
+    return experts.apply(
+        hidden_states=hidden_states_q,
+        w1=experts._w13_packed,
+        w2=experts._w2_packed,
+        router_logits=router_logits,
+        activation=MoEActivation.SILU,
+        global_num_experts=num_experts,
+        expert_map=None,
+        a1q_scale=hidden_states_scale,
+        apply_router_weight_on_input=False,
+        num_expert_group=_DSV3_N_GROUP,
+        topk_group=_DSV3_TOPK_GROUP,
+        e_score_correction_bias=routing_bias,
+        routed_scaling_factor=1.0,
+    )
+
+
+@pytest.mark.parametrize("num_tokens", [2, 7, 16])
+@pytest.mark.parametrize("top_k", [2, 4])
+def test_trtllm_nvfp4_monolithic_routing_replay_records_valid_experts(
+    num_tokens: int,
+    top_k: int,
+) -> None:
+    """End-to-end: ``TrtLlmNvFp4ExpertsMonolithic`` captures valid expert IDs
+    on the DSV3 routing path."""
+    if top_k > _DSV3_N_GROUP * _DSV3_TOPK_GROUP:
+        pytest.skip(
+            f"DSV3 requires top_k <= n_group * topk_group "
+            f"({_DSV3_N_GROUP * _DSV3_TOPK_GROUP})"
+        )
+    from flashinfer import fp4_quantize
+
+    torch.manual_seed(0)
+    device = torch.device("cuda:0")
+
+    num_experts = _DSV3_NUM_EXPERTS
+    hidden_size = 1024
+    intermediate_size = 1024
+    block_size = 16
+
+    experts, w13_q, _w13_s, w2_q, _w2_s, a_gs = _make_nvfp4_monolithic_experts(
+        num_experts=num_experts,
+        top_k=top_k,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        device=device,
+    )
+    # The apply() reads w1/w2 from its args, but we keep them on the experts
+    # for convenience of the helper.
+    experts._w13_packed = w13_q
+    experts._w2_packed = w2_q
+
+    assert experts.supports_routing_replay_capture()
+    captured: list[torch.Tensor] = []
+    experts.set_capture_fn(lambda r: captured.append(r.clone()))
+
+    hidden_states = (
+        torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) * 0.1
+    )
+    hidden_states_q, hidden_states_scale = fp4_quantize(
+        hidden_states,
+        a_gs,
+        block_size,
+        sf_use_ue8m0=False,
+        is_sf_swizzled_layout=False,
+    )
+    # The vLLM apply() does the .view(fp8_e4m3fn).reshape itself, so leave
+    # ``hidden_states_scale`` in its native (uint8 packed) form.
+    router_logits = torch.rand(
+        num_tokens, num_experts, device=device, dtype=torch.float32
+    )
+    routing_bias = _make_dsv3_routing_bias(num_experts, device)
+
+    _ = _run_nvfp4_monolithic(
+        experts,
+        hidden_states_q=hidden_states_q,
+        hidden_states_scale=hidden_states_scale,
+        router_logits=router_logits,
+        num_experts=num_experts,
+        routing_bias=routing_bias,
+    )
+
+    assert len(captured) == 1
+    replay = captured[0]
+    assert replay.dtype == torch.int16
+    assert replay.shape == (num_tokens, top_k)
+    assert (replay >= 0).all(), f"got out-of-range values: {replay}"
+    assert (replay < num_experts).all(), f"got out-of-range values: {replay}"
+    for t in range(num_tokens):
+        unique = replay[t].unique()
+        assert unique.numel() == top_k, (
+            f"token {t}: expected {top_k} distinct experts, "
+            f"got {unique.numel()} ({replay[t].tolist()})"
+        )
@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm import _custom_ops as ops
+
+
+def test_cutlass_group_gemm_python_guard_allows_thor(monkeypatch):
+    seen_capabilities: list[int] = []
+
+    def fake_cutlass_group_gemm_supported(capability: int) -> bool:
+        seen_capabilities.append(capability)
+        return True
+
+    monkeypatch.setattr(
+        torch.ops,
+        "_C",
+        SimpleNamespace(cutlass_group_gemm_supported=fake_cutlass_group_gemm_supported),
+    )
+
+    # CUDA 12 reports Thor as SM101 and CUDA 13 reports it as SM110. Both
+    # should reach the C++ query for the SM10x/SM11x CUTLASS MoE kernel.
+    assert ops.cutlass_group_gemm_supported(101)
+    assert ops.cutlass_group_gemm_supported(110)
+
+    # SM120 uses separate kernels and must not be advertised by this path.
+    assert not ops.cutlass_group_gemm_supported(120)
+    assert seen_capabilities == [101, 110]
@@ -0,0 +1,306 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import ast
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any, cast
+
+import pytest
+
+from vllm.model_executor.warmup.jit_warmup import (
+    VllmJitKernel,
+    WarmupIntRange,
+    get_ast_full_name,
+    zip_inputs,
+)
+
+
+def _next_power_of_2(value: int) -> int:
+    return 1 << max(0, value - 1).bit_length()
+
+
+def _round_up(value: int, *, multiple: int) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+def _config(
+    *,
+    bias: int = 0,
+    disabled: bool = False,
+    name: str = "base",
+    vectorized: bool = False,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        bias=bias,
+        disabled=disabled,
+        name=name,
+        vectorized=vectorized,
+    )
+
+
+class ToyKernel(VllmJitKernel["ToyKernel.CompileKey"]):
+    @dataclass(frozen=True)
+    class CompileKey:
+        block_size: int
+        work: int
+        vector_width: int
+        descriptor: tuple[object, ...]
+        enabled: bool
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        tokens: int,
+        cfg: Any,
+        lanes: int = 1,
+        mode: str = "default",
+        debug: int = 0,
+    ) -> CompileKey:
+        block_size = _next_power_of_2(tokens)
+        work: int = block_size * lanes + cfg.bias
+        return self.CompileKey(
+            block_size=block_size,
+            work=work,
+            vector_width=4 if cfg.vectorized and block_size >= 4 else 1,
+            descriptor=(
+                cfg.name,
+                mode,
+                -block_size,
+                block_size % 3,
+                block_size**2,
+            ),
+            enabled=not cfg.disabled,
+        )
+
+    def get_warmup_keys(self, max_tokens: int, cfg: Any) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            tokens=WarmupIntRange(1, max_tokens + 1),
+            cfg=cfg,
+            # This argument is intentionally unused by dispatch expressions.
+            debug=WarmupIntRange(0, 100),
+        )
+
+    def compile(self, compile_key: CompileKey) -> None:
+        pass
+
+
+class RecordingToyKernel(ToyKernel):
+    def __init__(self) -> None:
+        self.compiled: list[ToyKernel.CompileKey] = []
+        super().__init__()
+
+    def compile(self, compile_key: ToyKernel.CompileKey) -> None:
+        self.compiled.append(compile_key)
+
+
+def test_trace_dispatch_expands_ranges_dedupes_and_ignores_unused_inputs() -> None:
+    cfg = _config()
+
+    assert ToyKernel().get_warmup_keys(5, cfg) == [
+        ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
+        ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
+        ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
+        ToyKernel.CompileKey(8, 8, 1, ("base", "default", -8, 2, 64), True),
+    ]
+
+
+def test_compile_key_uses_defaults_locals_attributes_and_expressions() -> None:
+    cfg = _config(bias=3, disabled=True, name="cfg", vectorized=True)
+
+    assert ToyKernel().compile_key(
+        {
+            "tokens": 4,
+            "cfg": cfg,
+            "lanes": 2,
+        }
+    ) == ToyKernel.CompileKey(
+        block_size=4,
+        work=11,
+        vector_width=4,
+        descriptor=("cfg", "default", -4, 1, 16),
+        enabled=False,
+    )
+
+
+def test_trace_dispatch_combines_zipped_rows_with_independent_values() -> None:
+    cfg = _config(vectorized=True)
+
+    keys = ToyKernel()._trace_dispatch(ToyKernel().dispatch)(
+        zip_inputs(
+            dict(tokens=1, mode="small"),
+            dict(tokens=4, mode="wide"),
+        ),
+        cfg=cfg,
+        lanes=(1, 2),
+    )
+
+    assert keys == [
+        ToyKernel.CompileKey(1, 1, 1, ("base", "small", -1, 1, 1), True),
+        ToyKernel.CompileKey(1, 2, 1, ("base", "small", -1, 1, 1), True),
+        ToyKernel.CompileKey(4, 4, 4, ("base", "wide", -4, 1, 16), True),
+        ToyKernel.CompileKey(4, 8, 4, ("base", "wide", -4, 1, 16), True),
+    ]
+
+
+def test_zip_inputs_validates_input_rows() -> None:
+    with pytest.raises(ValueError, match="requires at least one"):
+        zip_inputs()
+    with pytest.raises(ValueError, match="rows must be mappings"):
+        zip_inputs(cast(Any, ("tokens", 1)))
+    with pytest.raises(ValueError, match="at least one dispatch input name"):
+        zip_inputs({})
+    with pytest.raises(ValueError, match="dispatch input names must be strings"):
+        zip_inputs(cast(Any, {1: 2}))
+    with pytest.raises(ValueError, match="same dispatch input names"):
+        zip_inputs({"tokens": 1}, {"mode": "small"})
+
+
+def test_trace_dispatch_rejects_bad_positional_groups_and_duplicates() -> None:
+    kernel = ToyKernel()
+
+    with pytest.raises(TypeError, match="zip_inputs"):
+        kernel._trace_dispatch(kernel.dispatch)(
+            cast(Any, {"tokens": 1}),
+            cfg=_config(),
+        )
+
+    with pytest.raises(ValueError, match="specified more than once"):
+        kernel._trace_dispatch(kernel.dispatch)(
+            zip_inputs(dict(tokens=1, mode="small")),
+            tokens=2,
+            cfg=_config(),
+        )
+
+
+def test_helper_calls_support_keywords_and_reject_star_kwargs() -> None:
+    class HelperKernel(VllmJitKernel["HelperKernel.CompileKey"]):
+        @dataclass(frozen=True)
+        class CompileKey:
+            value: int
+
+        def dispatch(  # type: ignore[override]
+            self,
+            *,
+            tokens: int,
+            block_size: int,
+        ) -> CompileKey:
+            return self.CompileKey(value=_round_up(tokens, multiple=block_size))
+
+        def get_warmup_keys(self) -> list[CompileKey]:
+            return []
+
+        def compile(self, compile_key: CompileKey) -> None:
+            pass
+
+    class StarKwargsKernel(VllmJitKernel["StarKwargsKernel.CompileKey"]):
+        @dataclass(frozen=True)
+        class CompileKey:
+            value: int
+
+        def dispatch(  # type: ignore[override]
+            self,
+            *,
+            tokens: int,
+            block_size: int,
+        ) -> CompileKey:
+            return self.CompileKey(value=_round_up(tokens, **{"multiple": block_size}))
+
+        def get_warmup_keys(self) -> list[CompileKey]:
+            return []
+
+        def compile(self, compile_key: CompileKey) -> None:
+            pass
+
+    assert HelperKernel().compile_key(
+        {
+            "tokens": 5,
+            "block_size": 4,
+        }
+    ) == HelperKernel.CompileKey(value=8)
+    with pytest.raises(ValueError, match=r"cannot use \*\*kwargs"):
+        StarKwargsKernel().compile_key({"tokens": 5, "block_size": 4})
+
+
+def test_dispatch_body_must_be_local_assignments_then_compile_key_return() -> None:
+    class BranchKernel(VllmJitKernel["BranchKernel.CompileKey"]):
+        @dataclass(frozen=True)
+        class CompileKey:
+            value: int
+
+        def dispatch(self, *, value: int) -> CompileKey:  # type: ignore[override]
+            if value > 0:
+                value = 1
+            return self.CompileKey(value=value)
+
+        def get_warmup_keys(self) -> list[CompileKey]:
+            return []
+
+        def compile(self, compile_key: CompileKey) -> None:
+            pass
+
+    class KwargsReturnKernel(VllmJitKernel["KwargsReturnKernel.CompileKey"]):
+        @dataclass(frozen=True)
+        class CompileKey:
+            value: int
+
+        def dispatch(self, *, value: int) -> CompileKey:  # type: ignore[override]
+            return self.CompileKey(**{"value": value})
+
+        def get_warmup_keys(self) -> list[CompileKey]:
+            return []
+
+        def compile(self, compile_key: CompileKey) -> None:
+            pass
+
+    with pytest.raises(ValueError, match="local assignments"):
+        BranchKernel()
+    with pytest.raises(ValueError, match=r"cannot use \*\*kwargs in CompileKey"):
+        KwargsReturnKernel()
+
+
+def test_dispatch_reports_unsupported_expression_with_context() -> None:
+    class UnsupportedKernel(VllmJitKernel["UnsupportedKernel.CompileKey"]):
+        @dataclass(frozen=True)
+        class CompileKey:
+            value: object
+
+        def dispatch(self, *, value: int) -> CompileKey:  # type: ignore[override]
+            return self.CompileKey(value={value})
+
+        def get_warmup_keys(self) -> list[CompileKey]:
+            return []
+
+        def compile(self, compile_key: CompileKey) -> None:
+            pass
+
+    with pytest.raises(ValueError) as exc_info:
+        UnsupportedKernel().compile_key({"value": 1})
+
+    message = str(exc_info.value)
+    assert "Unsupported dispatch expression" in message
+    assert "{value}" in message
+    assert "Supported dispatch expressions" in message
+
+
+def test_warmup_compiles_all_returned_keys_in_order() -> None:
+    kernel = RecordingToyKernel()
+    cfg = _config()
+
+    kernel.warmup(3, cfg)
+
+    assert kernel.compiled == [
+        ToyKernel.CompileKey(1, 1, 1, ("base", "default", -1, 1, 1), True),
+        ToyKernel.CompileKey(2, 2, 1, ("base", "default", -2, 2, 4), True),
+        ToyKernel.CompileKey(4, 4, 1, ("base", "default", -4, 1, 16), True),
+    ]
+
+
+def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
+    dotted_expr = ast.parse("foo.bar.baz").body[0]
+    call_expr = ast.parse("foo()").body[0]
+    assert isinstance(dotted_expr, ast.Expr)
+    assert isinstance(call_expr, ast.Expr)
+
+    assert get_ast_full_name(dotted_expr.value) == "foo.bar.baz"
+    assert get_ast_full_name(call_expr.value) is None
@@ -66,6 +66,12 @@ def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
    )


+def _make_modular_routed_experts():
+    return types.SimpleNamespace(
+        quant_method=types.SimpleNamespace(is_monolithic=False),
+    )
+
+
 def test_base_router_capture_pre_eplb_mapping():
    router = _make_router()
    captured = []
@@ -122,6 +128,8 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
        def __init__(self):
            self.layer_id = 7
            self.router = _make_router()
+            self.routed_experts = _make_modular_routed_experts()
+            self._quant_method = self.routed_experts.quant_method

    class DummyCapturer:
        def __init__(self):
@@ -160,6 +168,8 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
        def __init__(self):
            self.layer_id = 11
            self.router = _make_router()
+            self.routed_experts = _make_modular_routed_experts()
+            self._quant_method = self.routed_experts.quant_method

    class DummyCapturer:
        def __init__(self):
@@ -197,6 +207,8 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
        def __init__(self, layer_id):
            self.layer_id = layer_id
            self.router = _make_router()
+            self.routed_experts = _make_modular_routed_experts()
+            self._quant_method = self.routed_experts.quant_method

    target_module = DummyFusedMoE(layer_id=7)
    draft_module = DummyFusedMoE(layer_id=0)
@@ -222,6 +234,49 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
    assert draft_module.router.capture_fn is None


+def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
+    from vllm.v1.worker import gpu_model_runner as gmr
+
+    class DummyFusedMoE:
+        def __init__(self):
+            self.layer_id = 3
+            self.router = _make_router()
+            # Use a concrete monolithic expert and override its capability
+            # instead of instantiating the abstract base class directly.
+            from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
+                CPUExpertsFp8,
+            )
+
+            fused_experts = CPUExpertsFp8.__new__(CPUExpertsFp8)
+            self.routed_experts = types.SimpleNamespace(
+                quant_method=types.SimpleNamespace(
+                    is_monolithic=True,
+                    moe_kernel=types.SimpleNamespace(
+                        impl=types.SimpleNamespace(fused_experts=fused_experts)
+                    ),
+                )
+            )
+            self._quant_method = self.routed_experts.quant_method
+            self._quant_method.moe_kernel.impl.fused_experts = fused_experts
+            fused_experts.supports_routing_replay_capture = lambda: False
+
+    class DummyCapturer:
+        def capture(self, layer_id, topk_ids):
+            pass
+
+    dummy_module = DummyFusedMoE()
+    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
+
+    monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
+
+    dummy_self = types.SimpleNamespace(
+        model=types.SimpleNamespace(modules=lambda: [dummy_module])
+    )
+
+    with pytest.raises(ValueError, match="monolithic MoE kernel"):
+        gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
+
+
 def test_routed_experts_capturer_single_dp_no_metadata():
    """dp_metadata is None: capture writes the full topk_ids rows."""
    capturer = _capturer_with_buffer(dp_rank=0)
@@ -9,6 +9,7 @@ import torch.nn.functional as F
 from transformers import AutoModel

 from vllm.platforms import current_platform
+from vllm.utils.mem_constants import MiB_bytes

 from ....conftest import HfRunner
 from ....utils import VLLM_PATH
@@ -106,7 +107,12 @@ def test_prm_models(
    if current_platform.is_cpu():
        pytest.skip("CPU only supports V1")

-    with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
+    with vllm_runner(
+        model,
+        max_model_len=1024,
+        dtype=dtype,
+        kv_cache_memory_bytes=64 * MiB_bytes,
+    ) as vllm_model:
        vllm_outputs = vllm_model.token_classify(math_step_prompts)

    with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
@@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs(
    if not FIXTURE_REWARD_RESULT.get(model):
        pytest.skip(f"No available golden outputs for {model}.")

-    with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
+    with vllm_runner(
+        model,
+        max_model_len=1024,
+        dtype=dtype,
+        kv_cache_memory_bytes=64 * MiB_bytes,
+    ) as vllm_model:
        vllm_outputs = vllm_model.token_classify(math_step_prompts)

    golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
@@ -74,6 +74,44 @@ def test_cosmos3_new_checkpoint_weights_mapper():
    )


+def test_cosmos3_modelopt_quantizer_weights_mapper():
+    """ModelOpt/Diffusers FP8 checkpoints ship native fake-quant buffers
+    (``*_quantizer._amax`` / ``._scale``) alongside the vLLM-consumable
+    ``weight_scale`` / ``input_scale`` sidecars. vLLM must drop the former
+    (it has no parameter for them) while keeping the latter."""
+    from vllm.model_executor.models.cosmos3 import Cosmos3ForConditionalGeneration
+
+    mapper = Cosmos3ForConditionalGeneration.hf_to_vllm_mapper
+
+    # Native ModelOpt quantizer buffers are dropped.
+    assert (
+        mapper.apply_list(
+            [
+                "layers.0.self_attn.to_q.input_quantizer._amax",
+                "layers.0.self_attn.to_q.weight_quantizer._amax",
+                "layers.0.self_attn.to_q.weight_quantizer._scale",
+                "layers.0.mlp.down_proj.output_quantizer._amax",
+            ]
+        )
+        == []
+    )
+
+    # The FP8 scale sidecars vLLM actually consumes are kept and remapped.
+    assert mapper.apply_list(
+        [
+            "layers.0.self_attn.to_q.weight",
+            "layers.0.self_attn.to_q.weight_scale",
+            "layers.0.self_attn.to_q.input_scale",
+            "layers.0.mlp.down_proj.input_scale",
+        ]
+    ) == [
+        "language_model.model.layers.0.self_attn.q_proj.weight",
+        "language_model.model.layers.0.self_attn.q_proj.weight_scale",
+        "language_model.model.layers.0.self_attn.q_proj.input_scale",
+        "language_model.model.layers.0.mlp.down_proj.input_scale",
+    ]
+
+
 def test_cosmos3_edge_checkpoint_weights_mapper():
    from vllm.model_executor.models.cosmos3_edge import (
        Cosmos3EdgeForConditionalGeneration,
@@ -9,8 +9,8 @@ Note: these tests will only pass on L4 GPU.
 import pytest

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


@@ -70,13 +70,7 @@ def test_models(
    if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
        pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")

-    if not (
-        current_platform.is_xpu()
-        or (
-            get_flash_attn_version() == 3
-            and current_platform.is_device_capability_family(90)
-        )
-    ):
+    if not flash_attn_supports_kv_cache_dtype(kv_cache_dtype):
        pytest.skip(
            f"{kv_cache_dtype} is not supported on this GPU type with {backend} attention."
        )
@@ -6,6 +6,7 @@ import mimetypes
 import os
 import shutil
 import time
+from io import BytesIO
 from tempfile import NamedTemporaryFile, TemporaryDirectory

 import aiohttp
@@ -111,6 +112,34 @@ async def test_fetch_image_base64(
        assert _image_equals(data_image_sync, data_image_async)


+@pytest.mark.asyncio
+async def test_fetch_image_keep_original_mode():
+    """media_io_kwargs can disable the default RGB conversion."""
+    # RGBA image: opaque black pixel on a fully transparent background
+    rgba_image = Image.new("RGBA", (4, 4), (0, 0, 0, 0))
+    rgba_image.putpixel((2, 2), (0, 0, 0, 255))
+    buffer = BytesIO()
+    rgba_image.save(buffer, "PNG")
+    data_url = (
+        f"data:image/png;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+    )
+
+    # Default behavior: RGBA is composited onto a white background
+    default_image = MediaConnector().fetch_image(data_url)
+    assert default_image.mode == "RGB"
+    assert default_image.getpixel((0, 0)) == (255, 255, 255)
+    assert default_image.getpixel((2, 2)) == (0, 0, 0)
+
+    # image_mode=None via media_io_kwargs: original mode is preserved
+    connector = MediaConnector(media_io_kwargs={"image": {"image_mode": None}})
+    image_sync = connector.fetch_image(data_url)
+    image_async = await connector.fetch_image_async(data_url)
+    for image in (image_sync, image_async):
+        assert image.mode == "RGBA"
+        assert image.getpixel((0, 0)) == (0, 0, 0, 0)
+        assert image.getpixel((2, 2)) == (0, 0, 0, 255)
+
+
@pytest.mark.asyncio
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
 async def test_fetch_image_local_files(image_url: str):
@@ -80,6 +80,29 @@ def test_image_media_io_rgba_custom_background(tmp_path):
    assert green_numpy[0][0][2] == 0  # B


+def test_image_media_io_no_mode_conversion(tmp_path):
+    """image_mode=None skips conversion and preserves the original mode."""
+    # RGBA image: opaque black pixel on a fully transparent background
+    rgba_image = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
+    rgba_image.putpixel((5, 5), (0, 0, 0, 255))
+    test_image_path = tmp_path / "test_rgba.png"
+    rgba_image.save(test_image_path)
+
+    # Default behavior: RGBA is composited onto a white background
+    image_io_default = ImageMediaIO()
+    converted_default = image_io_default.load_file(test_image_path)
+    assert converted_default.media.mode == "RGB"
+    assert converted_default.media.getpixel((0, 0)) == (255, 255, 255)
+    assert converted_default.media.getpixel((5, 5)) == (0, 0, 0)
+
+    # image_mode=None: original mode and alpha channel are preserved
+    image_io_keep = ImageMediaIO(image_mode=None)
+    converted_keep = image_io_keep.load_file(test_image_path)
+    assert converted_keep.media.mode == "RGBA"
+    assert converted_keep.media.getpixel((0, 0)) == (0, 0, 0, 0)
+    assert converted_keep.media.getpixel((5, 5)) == (0, 0, 0, 255)
+
+
 def test_image_media_io_rgba_background_color_validation():
    """Test that invalid rgba_background_color values are properly rejected."""

@@ -165,6 +165,38 @@ def test_modelopt_mixed_precision_does_not_quantize_unlisted_fused_sibling():
    assert config._resolve_quant_algo("model.layers.0.linear_attn.in_proj_ba") is None


+def test_modelopt_mixed_precision_composes_gemma4_mappers():
+    from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+    from vllm.model_executor.models.gemma4_mm import (
+        Gemma4ForConditionalGeneration,
+    )
+
+    config = _mixed_precision_config(
+        {
+            "model.language_model.layers.0.experts": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+            "model.language_model.layers.1.moe.experts.gate_up_proj": {
+                "quant_algo": "NVFP4",
+                "group_size": 16,
+            },
+        }
+    )
+
+    config.apply_vllm_mapper(
+        Gemma4ForConditionalGeneration.hf_to_vllm_mapper.get_unstacked_mapper()
+    )
+    config.apply_vllm_mapper(Gemma4ForCausalLM.hf_to_vllm_mapper.get_unstacked_mapper())
+
+    expected_prefix = "language_model.model.layers.0.moe.experts"
+    assert set(config.quantized_layers) == {
+        expected_prefix,
+        "language_model.model.layers.1.moe.gate_up_proj",
+    }
+    assert config._resolve_quant_algo(expected_prefix) == "NVFP4"
+
+
 def test_modelopt_mixed_precision_infers_fused_gate_up_projection():
    from vllm.model_executor.layers.linear import LinearBase

@@ -67,6 +67,23 @@ def test_memory_profiling():
    non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024)  # noqa
    assert abs(non_torch_ratio - 1) <= 0.05
    assert result.torch_peak_increase == 1024 * 1024 * 1024
+
+    expected_total_consumed = (256 + 512) * 1024 * 1024
+    total_consumed_ratio = result.total_consumed / expected_total_consumed
+    assert abs(total_consumed_ratio - 1) <= 0.05, (
+        f"total_consumed={result.total_consumed}, "
+        f"expected={expected_total_consumed}, "
+        f"ratio={total_consumed_ratio}"
+    )
+
+    expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
+    non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
+    assert abs(non_kv_ratio - 1) <= 0.05, (
+        f"non_kv_cache_memory={result.non_kv_cache_memory}, "
+        f"expected={expected_non_kv}, "
+        f"ratio={non_kv_ratio}"
+    )
+
    del weights
    lib.cudaFree(handle1)
    lib.cudaFree(handle2)
@@ -21,6 +21,7 @@ from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import (
    STR_DTYPE_TO_TORCH_DTYPE,
+    is_quantized_kv_cache,
    is_torch_equal_or_newer,
    set_random_seed,
 )
@@ -45,6 +46,11 @@ BACKENDS_TO_TEST = [

 DEVICE_TYPE = current_platform.device_type

+FP8_KV_CACHE_DTYPES = {
+    "fp8": torch.float8_e4m3fn,
+    "fp8_e4m3": torch.float8_e4m3fn,
+}
+
 # Remove flashinfer from the list if it's not available
 try:
    import flashinfer  # noqa: F401
@@ -110,6 +116,7 @@ def create_and_prepopulate_kv_cache(
    num_blocks: int,
    common_attn_metadata: CommonAttentionMetadata,
    randomize_blocks: bool = True,
+    kv_cache_dtype: str = "auto",
 ) -> torch.Tensor:
    """Create and prepopulate a KV cache with context data.

@@ -140,8 +147,18 @@ def create_and_prepopulate_kv_cache(
    block_table = common_attn_metadata.block_table_tensor
    slot_mapping = common_attn_metadata.slot_mapping

+    # For an fp8 kv cache, store the cache in the fp8 dtype so that assigning
+    # the higher-precision context tensors quantizes them, mirroring runtime.
+    fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
+    storage_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype] if fp8_kv_cache else dtype
+
    kv_cache = torch.zeros(
-        num_blocks, block_size, num_kv_heads, 2 * head_size, dtype=dtype, device=device
+        num_blocks,
+        block_size,
+        num_kv_heads,
+        2 * head_size,
+        dtype=storage_dtype,
+        device=device,
    )
    kv_cache_flat = kv_cache.view(-1, num_kv_heads, 2 * head_size)

@@ -195,7 +212,12 @@ def create_and_prepopulate_kv_cache(
        ] * block_size + token_inter_block_offsets.to(device)

    # Transpose to logical (num_blocks, num_kv_heads, block_size, 2*hs)
-    return kv_cache.transpose(1, 2).contiguous()
+    kv_cache = kv_cache.transpose(1, 2).contiguous()
+
+    if fp8_kv_cache:
+        kv_cache = kv_cache.view(torch.uint8)
+
+    return kv_cache


 class MockAttentionLayer:
@@ -224,6 +246,7 @@ def run_attention_backend(
    kv_cache: torch.Tensor,
    attn_type: AttentionType = AttentionType.DECODER,
    sliding_window: int | None = None,
+    kv_cache_dtype: str = "auto",
 ) -> torch.Tensor:
    """Run attention computation using the specified backend's AttentionImpl."""

@@ -291,13 +314,16 @@ def run_attention_backend(
        alibi_slopes=None,
        sliding_window=sliding_window,
        attn_type=attn_type,
-        kv_cache_dtype="auto",
+        kv_cache_dtype=kv_cache_dtype,
    )

    # Create mock layer and output buffer
    mock_layer = MockAttentionLayer(device)
    output = torch.empty_like(query)

+    if is_quantized_kv_cache(kv_cache_dtype) and impl.supports_quant_query_input:
+        query = query.to(current_platform.fp8_dtype())
+
    # Run forward pass
    # NOTE: The query, key, and value are already shaped correctly
    # in the calling test function.
@@ -324,6 +350,7 @@ def _test_backend_correctness(
    atol: float = 1e-2,
    rtol: float = 1e-2,
    tensor_parallel_size: int = 1,
+    kv_cache_dtype: str = "auto",
 ):
    """
    Test that all backends produce similar outputs to a reference implementation
@@ -372,6 +399,7 @@ def _test_backend_correctness(
        num_gpu_blocks=8192,
        hf_config_override=hf_config_override,
    )
+    vllm_config.cache_config.cache_dtype = kv_cache_dtype
    device = torch.device(f"{DEVICE_TYPE}:0")

    kv_cache_spec = create_standard_kv_cache_spec(vllm_config, attn_type)
@@ -392,6 +420,13 @@ def _test_backend_correctness(
    block_size = vllm_config.cache_config.block_size
    scale = 1.0 / (head_size**0.5)

+    fp8_kv_cache = is_quantized_kv_cache(kv_cache_dtype)
+    if fp8_kv_cache:
+        query_fp8_dtype = current_platform.fp8_dtype()
+        kv_fp8_dtype = FP8_KV_CACHE_DTYPES[kv_cache_dtype]
+        atol = max(atol, 6e-2)
+        rtol = max(rtol, 1e-1)
+
    # 2. Generate data and compute SDPA reference output
    all_q_vllm, all_k_vllm, all_v_vllm = [], [], []
    all_sdpa_outputs = []
@@ -407,10 +442,17 @@ def _test_backend_correctness(
        k_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)
        v_full = torch.randn(s_len, num_kv_heads, head_size, dtype=dtype, device=device)

+        if fp8_kv_cache:
+            q_ref = q.to(query_fp8_dtype).to(dtype)
+            k_ref = k_full.to(kv_fp8_dtype).to(dtype)
+            v_ref = v_full.to(kv_fp8_dtype).to(dtype)
+        else:
+            q_ref, k_ref, v_ref = q, k_full, v_full
+
        # SDPA expects (N, H, L, D), so unsqueeze batch and permute
-        q_sdpa_in = q.unsqueeze(0).transpose(1, 2)
-        k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
-        v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
+        q_sdpa_in = q_ref.unsqueeze(0).transpose(1, 2)
+        k_sdpa_in = k_ref.unsqueeze(0).transpose(1, 2)
+        v_sdpa_in = v_ref.unsqueeze(0).transpose(1, 2)

        if num_q_heads != num_kv_heads:
            assert num_q_heads % num_kv_heads == 0, (
@@ -471,6 +513,7 @@ def _test_backend_correctness(
        num_blocks=vllm_config.cache_config.num_gpu_blocks or 1000,
        common_attn_metadata=common_attn_metadata,
        randomize_blocks=True,
+        kv_cache_dtype=kv_cache_dtype,
    )

    # 4. Run vLLM backends and compare
@@ -488,6 +531,12 @@ def _test_backend_correctness(
        else:
            backend_cls = None

+        if is_quantized_kv_cache(kv_cache_dtype) and (
+            backend_cls is None
+            or not backend_cls.supports_kv_cache_dtype(kv_cache_dtype)
+        ):
+            continue
+
        if backend_name == AttentionBackendEnum.FLASHINFER:
            set_kv_cache_layout("HND")
            reset_kv_cache_layout = True
@@ -521,6 +570,7 @@ def _test_backend_correctness(
                kv_cache_for_backend,
                sliding_window=sliding_window,
                attn_type=attn_type,
+                kv_cache_dtype=kv_cache_dtype,
            )
        finally:
            if reset_kv_cache_layout:
@@ -570,8 +620,13 @@ def _test_backend_correctness(
 )
@pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
 def test_causal_backend_correctness(
-    default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int
+    default_vllm_config,
+    batch_spec_name: str,
+    model: str,
+    tensor_parallel_size: int,
+    kv_cache_dtype: str,
 ):
    """Test backend's correctness with causal attention."""

@@ -612,6 +667,7 @@ def test_causal_backend_correctness(
        SMALL_BLOCK_BACKENDS,
        causal_mask_mod,
        tensor_parallel_size=tensor_parallel_size,
+        kv_cache_dtype=kv_cache_dtype,
    )

    # Fast FlexAttention needs to run with block_size=128
@@ -623,6 +679,7 @@ def test_causal_backend_correctness(
            causal_mask_mod,
            block_size=128,
            tensor_parallel_size=tensor_parallel_size,
+            kv_cache_dtype=kv_cache_dtype,
        )


@@ -119,6 +119,7 @@ def test_mla_post_load_preserves_runtime_weight_addresses(monkeypatch):
    layer.kv_b_proj.quant_method = None
    layer.is_aiter_triton_fp4_bmm_enabled = False
    layer.is_aiter_triton_fp8_bmm_enabled = False
+    layer.dcp_q_replicate = False
    layer.quant_config = None
    layer.layer_name = "test"

@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (
    EncoderOnlyAttentionSpec,
    FullAttentionSpec,
    MambaSpec,
+    get_kv_quant_mode,
 )


@@ -178,6 +179,7 @@ def create_standard_kv_cache_spec(
        head_size=vllm_config.model_config.get_head_size(),
        dtype=vllm_config.model_config.dtype,
        sliding_window=vllm_config.model_config.get_sliding_window(),
+        kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype),
    )


@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import copy
 import hashlib
 import importlib
 from collections.abc import Callable
@@ -1477,6 +1478,128 @@ def test_get_max_concurrency_for_kv_cache_config():
    assert num_tokens == max_concurrency_hybrid_model * max_model_len
    assert max_concurrency == max_concurrency_hybrid_model

+    # Unequal group sizes in the standard layout: each group's pages cost
+    # whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
+    # same as the equal-hybrid case above, regardless of the second group
+    # holding only 2 layers.
+    kv_cache_config_unequal_groups = KVCacheConfig(
+        num_blocks=1153 * 3,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
+            KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec),
+        ],
+    )
+    assert (
+        get_max_concurrency_for_kv_cache_config(
+            vllm_config, kv_cache_config_unequal_groups
+        )
+        == 3
+    )
+
+    # UniformTypeKVCacheSpecs group (worker config shape): the aggregated
+    # spec's memory/page ratio equals a single layer's page count, so the
+    # group needs 1024 blocks and the request 1153 in total. The previous
+    # formula normalized both groups' memory by the first group's page size,
+    # reporting 3459/1057 = 3.27 here instead of 3 — and a different value
+    # again for the scheduler-config shape below.
+    uniform_full_spec = UniformTypeKVCacheSpecs(
+        block_size=full_attention_spec.block_size,
+        kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
+    )
+    kv_cache_config_uniform_group = KVCacheConfig(
+        num_blocks=1153 * 3,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),
+            KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec),
+        ],
+    )
+    assert (
+        get_max_concurrency_for_kv_cache_config(
+            vllm_config, kv_cache_config_uniform_group
+        )
+        == 3
+    )
+
+    # Scheduler-config shape: generate_scheduler_kv_cache_config replaces the
+    # uniform-type group's spec with a representative per-layer spec.
+    # Capacity must not change between the two shapes (the engine computes
+    # on the scheduler config, the worker loop on the worker config).
+    kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config(
+        [copy.deepcopy(kv_cache_config_uniform_group)]
+    )
+    assert get_max_concurrency_for_kv_cache_config(
+        vllm_config, kv_cache_config_scheduler_shape
+    ) == get_max_concurrency_for_kv_cache_config(
+        vllm_config, kv_cache_config_uniform_group
+    )
+
+
+def test_get_max_concurrency_packed_kv_cache_config():
+    from vllm.v1.core.kv_cache_utils import (
+        _get_kv_cache_config_packed,
+        _use_packed_kv_cache_config,
+    )
+
+    model_config = ModelConfig(
+        "Qwen/Qwen1.5-7B",
+        runner="generate",
+        dtype="float16",
+        max_model_len=16384,
+    )
+    scheduler_config = SchedulerConfig(
+        max_num_batched_tokens=1024,
+        enable_chunked_prefill=True,
+        max_model_len=model_config.max_model_len,
+        is_encoder_decoder=model_config.is_encoder_decoder,
+        async_scheduling=False,
+    )
+    vllm_config = VllmConfig(
+        model_config=model_config,
+        scheduler_config=scheduler_config,
+    )
+
+    # All-UniformTypeKVCacheSpecs groups select the packed layout.
+    mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
+    swa_specs = {
+        f"layer_{i}": SlidingWindowMLASpec(
+            block_size=16,
+            num_kv_heads=1,
+            head_size=576,
+            dtype=torch.float32,
+            sliding_window=128,
+        )
+        for i in range(4, 6)
+    }
+    kv_cache_groups = [
+        KVCacheGroupSpec(
+            list(mla_specs),
+            UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
+        ),
+        KVCacheGroupSpec(
+            list(swa_specs),
+            UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
+        ),
+    ]
+    assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
+    num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
+        vllm_config, kv_cache_groups, 2 * GiB_bytes
+    )
+    assert num_blocks > 0
+    kv_cache_config_packed = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=kv_cache_tensors,
+        kv_cache_groups=kv_cache_groups,
+    )
+    # Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
+    # the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
+    # previous formula normalized by the first group's page size and gave
+    # 1061 blocks per request instead of 1097.
+    assert get_max_concurrency_for_kv_cache_config(
+        vllm_config, kv_cache_config_packed
+    ) == num_blocks / (1024 + 73)
+

 def test_allocate_with_lookahead():
    """Verify that lookahead tokens correctly affect block allocation"""
@@ -7,7 +7,13 @@ import torch

 from tests.v1.kv_connector.unit.utils import create_vllm_config
 from vllm.config import KVEventsConfig, KVTransferConfig
-from vllm.distributed.kv_events import MEDIUM_CPU, MEDIUM_FS, BlockRemoved, BlockStored
+from vllm.distributed.kv_events import (
+    MEDIUM_CPU,
+    MEDIUM_FS,
+    MEDIUM_OBJ,
+    BlockRemoved,
+    BlockStored,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
    build_offloading_config,
 )
@@ -62,8 +68,9 @@ def _wire_hash(block_hash: BlockHash):
    return maybe_convert_block_hash(block_hash)


-def _request(*, block_hashes: list[BlockHash], token_count: int):
+def _request(*, block_hashes: list[BlockHash], token_count: int, req_id: str = "req"):
    req = MagicMock()
+    req.request_id = req_id
    req.block_hashes = block_hashes
    req.all_token_ids = list(range(1, token_count + 1))
    req.lora_request = None
@@ -104,10 +111,32 @@ def _record_chunks(
    return keys


+def _record_lookup_chunks(
+    tracker: OffloadingEventsTracker,
+    req,
+    group_config: GroupOffloadConfig,
+    num_chunks: int,
+) -> list[OffloadKey]:
+    keys: list[OffloadKey] = []
+    hbf = group_config.hashes_per_chunk
+    for chunk_idx in range(num_chunks):
+        tail_hash = req.block_hashes[(chunk_idx + 1) * hbf - 1]
+        assert tail_hash is not None
+        key = make_offload_key(tail_hash, group_config.group_idx)
+        tracker.record_lookup(
+            req,
+            group_config,
+            chunk_idx,
+            key,
+        )
+        keys.append(key)
+    return keys
+
+
 def _stored_event(
    keys: list[OffloadKey],
-    locality: Locality | None = None,
    medium: str = _CPU_MEDIUM,
+    locality: Locality | None = None,
 ) -> OffloadingEvent:
    return OffloadingEvent(
        keys=keys,
@@ -119,8 +148,8 @@ def _stored_event(

 def _removed_event(
    keys: list[OffloadKey],
-    locality: Locality | None = None,
    medium: str = _CPU_MEDIUM,
+    locality: Locality | None = None,
 ) -> OffloadingEvent:
    return OffloadingEvent(
        keys=keys,
@@ -130,6 +159,21 @@ def _removed_event(
    )


+def _lookup_chunk() -> tuple[
+    OffloadingEventsTracker, MagicMock, GroupOffloadConfig, OffloadKey
+]:
+    tracker = _tracker()
+    req = _request(block_hashes=[_hash(0)], token_count=4)
+    group_config = _group_config()
+    key = _record_lookup_chunks(
+        tracker,
+        req,
+        group_config,
+        num_chunks=1,
+    )[0]
+    return tracker, req, group_config, key
+
+
 def test_take_events_forwards_locality_to_rich_store():
    tracker = _tracker()
    req = _request(block_hashes=[_hash(0)], token_count=4)
@@ -220,18 +264,37 @@ def test_take_events_publishes_routable_block_stored():
    assert len(tracker._pending_event_metadata) == 6


-def test_take_events_factor_gt_1_chunk_store_and_remove():
+def test_promotion_emits_full_cpu_stored_event():
+    tracker, _, _, key = _lookup_chunk()
+
+    [event] = tracker.take_events([_stored_event([key])])
+
+    assert isinstance(event, BlockStored)
+    assert event.medium == MEDIUM_CPU
+    assert event.block_hashes == [_wire_hash(_hash(0))]
+    assert event.parent_block_hash is None
+    assert event.token_ids == [1, 2, 3, 4]
+    assert event.block_size == 4
+    assert event.lora_id is None
+    assert event.lora_name is None
+    assert event.extra_keys is None
+    assert event.group_idx == 0
+    assert event.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
+    assert event.kv_cache_spec_sliding_window is None
+
+
+def test_lookup_promotion_factor_gt_1_store_and_remove():
    block_size = 4
-    blocks_per_chunk = 3
+    blocks_per_chunk = 2
    tracker = _tracker()
    group_config = _group_config(
        block_size=block_size, blocks_per_chunk=blocks_per_chunk
    )
    req = _request(
-        block_hashes=[_hash(i) for i in range(6)],
+        block_hashes=[_hash(i) for i in range(4)],
        token_count=block_size * blocks_per_chunk * 2,
    )
-    keys = _record_chunks(tracker, req, group_config, num_chunks=2)
+    keys = _record_lookup_chunks(tracker, req, group_config, num_chunks=2)

    stored = list(tracker.take_events([_stored_event(keys)]))
    assert len(stored) == 2
@@ -293,6 +356,7 @@ def test_take_events_opt_out_keeps_placeholders():
    group_config = _group_config()
    req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
    keys = _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)

    assert not tracker.self_describing_enabled
    assert not tracker._pending_event_metadata
@@ -315,11 +379,21 @@ def test_take_events_opt_out_keeps_placeholders():
    assert len(events[3].block_hashes) == 3


-def test_record_store_skips_sliding_window_group():
+@pytest.mark.parametrize(
+    "sliding_window_size_in_chunks",
+    [1, 2],
+    ids=["ssm", "sliding-window"],
+)
+def test_event_metadata_skips_non_full_attention_group(
+    sliding_window_size_in_chunks: int,
+):
    tracker = _tracker()
-    group_config = _group_config(sliding_window_size_in_chunks=2)
+    group_config = _group_config(
+        sliding_window_size_in_chunks=sliding_window_size_in_chunks
+    )
    req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
    keys = _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)

    assert not tracker._pending_event_metadata

@@ -329,6 +403,57 @@ def test_record_store_skips_sliding_window_group():
    assert events[0].block_size == 0


+def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():
+    tracker = _tracker()
+    block_hashes = [_hash(0), _hash(1)]
+    req = _request(block_hashes=block_hashes, token_count=8)
+    group_config = _group_config(blocks_per_chunk=2)
+    key = _record_chunks(tracker, req, group_config, num_chunks=1)[0]
+    confirmed_meta = tracker._pending_event_metadata[key]
+    lookup_req = _request(
+        block_hashes=block_hashes,
+        token_count=8,
+        req_id="new-request",
+    )
+
+    tracker.record_lookup(
+        lookup_req,
+        group_config,
+        0,
+        key,
+    )
+    assert tracker._pending_event_metadata[key] is confirmed_meta
+
+    removed = list(tracker.take_events([_removed_event([key])]))
+    assert len(removed) == 1
+    assert removed[0].block_hashes == [
+        _wire_hash(_hash(0)),
+        _wire_hash(_hash(1)),
+    ]
+
+    stored = list(tracker.take_events([_stored_event([key])]))
+    assert len(stored) == 1
+    assert stored[0].block_size == 0
+    assert stored[0].token_ids == []
+
+    tracker.record_lookup(lookup_req, group_config, 0, key)
+    removed = list(tracker.take_events([_removed_event([key])]))
+    assert removed[0].block_hashes == [
+        _wire_hash(_hash(0)),
+        _wire_hash(_hash(1)),
+    ]
+
+
+@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ])
+def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str):
+    tracker, _, _, key = _lookup_chunk()
+    expected_metadata = dict(tracker._pending_event_metadata)
+
+    stored = list(tracker.take_events([_stored_event([key], medium)]))
+    assert stored[0].token_ids == [1, 2, 3, 4]
+    assert tracker._pending_event_metadata == expected_metadata
+
+
 def test_take_events_groups_removed_hashes_by_kv_group():
    tracker = _tracker()
    group0_config = _group_config(group_idx=0, blocks_per_chunk=2)
@@ -378,7 +503,7 @@ def test_reset_cache_clears_side_table():
    tracker = _tracker()
    group_config = _group_config()
    req = _request(block_hashes=[_hash(i) for i in range(3)], token_count=12)
-    _record_chunks(tracker, req, group_config, num_chunks=3)
+    _record_lookup_chunks(tracker, req, group_config, num_chunks=3)

    assert tracker._pending_event_metadata

@@ -387,7 +512,7 @@ def test_reset_cache_clears_side_table():
    assert not tracker._pending_event_metadata


-def test_tiering_rejects_self_describing_kv_events():
+def test_tiering_accepts_self_describing_kv_events():
    vllm_config = create_vllm_config(
        block_size=4,
        max_num_batched_tokens=16,
@@ -423,5 +548,9 @@ def test_tiering_rejects_self_describing_kv_events():
        ],
    )

-    with pytest.raises(ValueError, match="TieringOffloadingSpec"):
-        TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+    spec = TieringOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))
+    tracker = OffloadingEventsTracker(spec.kv_events_config)
+
+    assert spec.kv_events_config.enable_kv_cache_events
+    assert spec.kv_events_config.self_describing_kv_events
+    assert tracker.self_describing_enabled
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call

 import pytest
 import torch
@@ -11,6 +11,7 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (
    to_keys,
 )
 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
+from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
    OffloadingConnectorStats,
    _ConnectorMetricName,
@@ -18,7 +19,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
    OffloadingConnectorScheduler,
    RequestOffloadState,
+    is_store_reachable_swa_chunk,
 )
+from vllm.v1.core.kv_cache_utils import BlockHash
 from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheGroupSpec,
@@ -26,6 +29,7 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.kv_offload.base import (
    LookupResult,
+    OffloadingEvent,
    OffloadingManager,
    OffloadPolicy,
    ReqContext,
@@ -123,6 +127,55 @@ def test_scheduler_reports_lookup_sync_delay(request_runner):
    assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0


+@pytest.mark.parametrize(
+    (
+        "absolute_chunk_index",
+        "storable_chunk_count",
+        "alignment_chunk_count",
+        "sliding_window_chunks",
+        "is_eagle_group",
+        "expected",
+    ),
+    [
+        # Full 64-chunk segment: ordinary SWA keeps 62-63; EAGLE also keeps 61.
+        (61, 64, 64, 2, False, False),
+        (62, 64, 64, 2, False, True),
+        (60, 64, 64, 2, True, False),
+        (61, 64, 64, 2, True, True),
+        # Partial 48-of-64 segment: the reachable tail ends at chunk 47.
+        (45, 48, 64, 2, False, False),
+        (46, 48, 64, 2, False, True),
+        (44, 48, 64, 2, True, False),
+        (45, 48, 64, 2, True, True),
+        # A later partial segment uses its own actual end (chunks 64-79).
+        (76, 80, 64, 3, False, False),
+        (77, 80, 64, 3, False, True),
+        # No alignment means no store-pruning optimization.
+        (0, 1, None, None, False, True),
+        # A tail at least as large as the segment keeps every chunk.
+        (0, 2, 64, 2, False, True),
+    ],
+)
+def test_is_store_reachable_swa_chunk(
+    absolute_chunk_index: int,
+    storable_chunk_count: int,
+    alignment_chunk_count: int | None,
+    sliding_window_chunks: int | None,
+    is_eagle_group: bool,
+    expected: bool,
+):
+    assert (
+        is_store_reachable_swa_chunk(
+            absolute_chunk_index,
+            storable_chunk_count,
+            alignment_chunk_count,
+            sliding_window_chunks,
+            is_eagle_group,
+        )
+        is expected
+    )
+
+
 def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
    """A deferred lookup reports its async delay once it resolves."""
    runner = request_runner(
@@ -143,6 +196,159 @@ def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):
    assert reduced[f"{_ConnectorMetricName.LOOKUP_ASYNC_DELAY}_sum"] > 0


+def test_max_offload_tokens_zero_does_not_record_pending_lookups(request_runner):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+    )
+    runner.manager.lookup.return_value = LookupResult.RETRY
+    runner.manager.take_events.return_value = []
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    runner.new_request(
+        token_ids=[1] * 12,
+        kv_transfer_params={"max_offload_tokens": 0},
+    )
+    runner.run(decoded_tokens=[])
+
+    tracker = runner.connector_scheduler._events_tracker
+    assert runner.manager.lookup.call_count == 3
+    assert not tracker._pending_event_metadata
+    assert list(runner.connector_scheduler.take_events()) == []
+
+    runner.manager.lookup.return_value = LookupResult.MISS
+    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+
+    assert not tracker._pending_event_metadata
+    assert list(runner.connector_scheduler.take_events()) == []
+
+
+def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal(
+    request_runner,
+):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+        blocks_per_chunk=2,
+    )
+    raw_events: list[OffloadingEvent] = []
+
+    def take_raw_events():
+        yield from raw_events
+        raw_events.clear()
+
+    runner.manager.lookup.return_value = LookupResult.RETRY
+    runner.manager.take_events.side_effect = take_raw_events
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+
+    runner.new_request(token_ids=[1] * 8)
+    runner.run(decoded_tokens=[])
+
+    tracker = runner.connector_scheduler._events_tracker
+    assert not tracker._pending_event_metadata
+    key = runner.manager.lookup.call_args.args[0]
+    req_id = str(runner.req_id)
+    req_status = runner.connector_scheduler._req_status[req_id]
+
+    runner.scheduler.finish_requests((req_id,), RequestStatus.FINISHED_ABORTED)
+
+    assert not tracker._pending_event_metadata
+
+    raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
+    events = list(runner.connector_scheduler.take_events())
+    assert len(events) == 1
+    assert isinstance(events[0], BlockStored)
+    assert events[0].block_size == 0
+    assert events[0].token_ids == []
+
+    runner.manager.lookup.return_value = LookupResult.HIT
+    group_config = runner.connector_scheduler.config.kv_group_configs[0]
+    assert (
+        runner.connector_scheduler._maximal_prefix_lookup(
+            [key],
+            req_status.req_context,
+            req_status.req,
+            group_config,
+            0,
+        )
+        == 1
+    )
+    assert key in tracker._pending_event_metadata
+
+    raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=True))
+    [event] = runner.connector_scheduler.take_events()
+    assert isinstance(event, BlockRemoved)
+    assert event.medium == MEDIUM_CPU
+    assert len(event.block_hashes) == 2
+    assert key not in tracker._pending_event_metadata
+
+
+@pytest.mark.parametrize("blocks_per_chunk", [1, 2])
+def test_promotion_hit_precedes_stored_event_translation(
+    request_runner,
+    blocks_per_chunk: int,
+):
+    runner = request_runner(
+        block_size=4,
+        num_gpu_blocks=10,
+        async_scheduling=False,
+        blocks_per_chunk=blocks_per_chunk,
+    )
+    token_ids = [1] * 4 * blocks_per_chunk
+
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+    runner.new_request(token_ids=token_ids)
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=tuple(range(blocks_per_chunk)),
+    )
+    runner.scheduler.reset_prefix_cache()
+    runner.connector_scheduler._events_tracker.reset()
+
+    raw_events: list[OffloadingEvent] = []
+
+    def lookup(key, req_context):
+        raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))
+        return LookupResult.HIT
+
+    def take_raw_events():
+        yield from raw_events
+        raw_events.clear()
+
+    runner.manager.lookup.side_effect = lookup
+    runner.manager.take_events.side_effect = take_raw_events
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output([])
+    )
+    publisher = MagicMock()
+    runner.scheduler.kv_event_publisher = publisher
+
+    runner.new_request(token_ids=token_ids)
+    runner.run(
+        decoded_tokens=[],
+        expected_loaded=tuple(range(blocks_per_chunk)),
+    )
+
+    events = [
+        event
+        for publish_call in publisher.publish.call_args_list
+        for event in publish_call.args[0].events
+        if isinstance(event, BlockStored) and event.medium == MEDIUM_CPU
+    ]
+    assert len(events) == 1
+    assert len(events[0].block_hashes) == blocks_per_chunk
+    assert events[0].block_size == 4
+    assert events[0].token_ids == token_ids
+
+
@pytest.mark.parametrize("async_scheduling", [True, False])
 def test_offloading_connector(request_runner, async_scheduling: bool):
    block_size = 4
@@ -241,7 +447,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output([])
    )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2))

    # single block lookup with a hit in a middle block
@@ -249,7 +455,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output([])
    )
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5))


@@ -307,7 +513,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):

    # request should now return from preemption
    # re-load [0, ..., 8] from the CPU and store [9, 10, 11]
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output(keys)
    )
@@ -427,7 +633,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:
    # start a request to load the first block, but don't complete
    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.run(
        decoded_tokens=[],
        complete_transfers=False,
@@ -439,7 +645,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:

    # start a new request to load the same first block
    runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.run(
        decoded_tokens=[],
        complete_transfers=False,
@@ -491,7 +697,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):
    # start a request to load the first block, but don't complete
    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.run(
        decoded_tokens=[],
        complete_transfers=False,
@@ -793,73 +999,144 @@ def _make_scheduler_with_lookup(

    scheduler = object.__new__(OffloadingConnectorScheduler)
    scheduler.manager = manager
+    scheduler._events_tracker = MagicMock()
    return scheduler


 _EMPTY_REQ_CTX = ReqContext(req_id="")
+_LOOKUP_REQ = MagicMock()
+_LOOKUP_REQ.request_id = "req"
+_LOOKUP_GROUP_CONFIG = MagicMock()
+
+
+def _maximal_lookup(sched, keys, start_chunk_idx: int = 0):
+    return sched._maximal_prefix_lookup(
+        keys,
+        _EMPTY_REQ_CTX,
+        _LOOKUP_REQ,
+        _LOOKUP_GROUP_CONFIG,
+        start_chunk_idx,
+    )


 class TestMaximalPrefixLookup:
    def test_all_hit(self):
        sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 2
+
+    def test_records_absolute_chunk_indices(self):
+        keys = to_keys([1, 2])
+        sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
+
+        assert _maximal_lookup(sched, keys, start_chunk_idx=3) == 2
+        assert sched._events_tracker.record_lookup.call_args_list == [
+            call(
+                _LOOKUP_REQ,
+                _LOOKUP_GROUP_CONFIG,
+                3,
+                keys[0],
+            ),
+            call(
+                _LOOKUP_REQ,
+                _LOOKUP_GROUP_CONFIG,
+                4,
+                keys[1],
+            ),
+        ]

    def test_all_miss(self):
        sched = _make_scheduler_with_lookup({})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 0
+        sched._events_tracker.record_lookup.assert_not_called()

    def test_partial_prefix(self):
        sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) == 2

    def test_miss_then_hit(self):
        sched = _make_scheduler_with_lookup({2: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, to_keys([1, 2])) == 0

    def test_single_hit(self):
        sched = _make_scheduler_with_lookup({1: LookupResult.HIT})
-        assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1
+        assert _maximal_lookup(sched, to_keys([1])) == 1

    def test_empty(self):
        sched = _make_scheduler_with_lookup({})
-        assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0
+        assert _maximal_lookup(sched, []) == 0
+
+    @pytest.mark.parametrize(
+        "pending_result",
+        [LookupResult.RETRY, LookupResult.HIT_PENDING],
+    )
+    def test_pending_result_is_not_recorded(
+        self,
+        pending_result: LookupResult,
+    ):
+        sched = _make_scheduler_with_lookup({1: pending_result})
+
+        assert _maximal_lookup(sched, to_keys([1])) is None
+        sched._events_tracker.record_lookup.assert_not_called()

    def test_retry_defers(self):
+        keys = to_keys([1, 2])
        sched = _make_scheduler_with_lookup(
            {1: LookupResult.RETRY, 2: LookupResult.HIT}
        )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
        assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            1,
+            keys[1],
+        )

    def test_retry_after_hit_defers(self):
+        keys = to_keys([1, 2])
        sched = _make_scheduler_with_lookup(
            {1: LookupResult.HIT, 2: LookupResult.RETRY}
        )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            0,
+            keys[0],
+        )

    def test_hit_pending_defers(self):
+        keys = to_keys([1, 2])
        sched = _make_scheduler_with_lookup(
            {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT}
        )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, keys) is None
        assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_called_once_with(
+            _LOOKUP_REQ,
+            _LOOKUP_GROUP_CONFIG,
+            1,
+            keys[1],
+        )

    def test_hit_pending_does_not_stop_scan(self):
        """HIT_PENDING defers but does not break — scan continues until miss."""
        sched = _make_scheduler_with_lookup(
            {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT}
        )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
        assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_not_called()

    def test_retry_stops_at_miss(self):
        """RETRY is treated as hit for iteration, but miss stops the scan."""
        sched = _make_scheduler_with_lookup(
            {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT}
        )
-        assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None
+        assert _maximal_lookup(sched, to_keys([1, 2, 3])) is None
        # lookup should have been called for blocks 1 and 2 (stops at miss)
        assert sched.manager.lookup.call_count == 2
+        sched._events_tracker.record_lookup.assert_not_called()


 class TestSlidingWindowLookup:
@@ -1011,7 +1288,7 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling

    # New request with 2 offloaded chunks; first matches what's in CPU.
    runner.new_request(token_ids=[0] * tokens_per_chunk * 2)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output(keys)
    )
@@ -1042,7 +1319,7 @@ def test_loads_do_not_populate_fence_index(request_runner):
        async_scheduling=False,
    )
    runner.new_request(token_ids=[0] * 12)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.run(decoded_tokens=[], complete_transfers=False)
    assert runner.connector_scheduler._block_id_to_pending_jobs == {}

@@ -1088,7 +1365,7 @@ def test_fence_at_update_state_after_alloc(request_runner):

    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[0] * 4)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output([])
    )
@@ -1139,7 +1416,7 @@ def test_fence_at_build_store_jobs(request_runner):

    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[1] * 4)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 0
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 0
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output([])
    )
@@ -1365,7 +1642,7 @@ def test_reset_cache(request_runner, async_scheduling: bool):
    # Leave the load in-flight so that reset_cache must flush it.
    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[0] * tokens_per_chunk)
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1
    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
        generate_store_output([])
    )
@@ -1553,9 +1830,7 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner
    # 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
-    )
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: len(keys)

    readmit_output = runner.scheduler.schedule()

@@ -1669,7 +1944,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):
    runner.scheduler.reset_prefix_cache()
    runner.new_request(token_ids=[0] * num_tokens + [1])
    runner.manager.lookup.return_value = LookupResult.HIT
-    runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2
+    runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2
    runner.run(
        decoded_tokens=[EOS_TOKEN_ID],
        # Group 0: full prefix lookup hits 2 offloaded chunks
@@ -1840,6 +2115,13 @@ class TestEagle:
        req.request_id = "test-req"
        req.num_tokens = num_tokens
        req.kv_transfer_params = None
+        num_hash_blocks = max(
+            len(hashes) * scheduler.config.kv_group_configs[idx].hashes_per_chunk
+            for idx, hashes in enumerate(offload_keys_per_group)
+        )
+        req.block_hashes = [BlockHash(str(i).encode()) for i in range(num_hash_blocks)]
+        req.all_token_ids = list(range(num_tokens))
+        req.lora_request = None

        state = RequestOffloadState(
            config=scheduler.config,
@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the EAGLE speculator's draft attention metadata builder.
+
+These tests guard the regression where ``_build_draft_attn_metadata`` did
+not populate ``seq_lens_cpu_upper_bound`` on the per-step
+``CommonAttentionMetadata``. Several downstream attention backends and
+helpers (``split_decodes_prefills_and_extends``, the MLA indexer,
+flex-attention, cross-attention) assert this field is non-None, so
+omitting it caused crashes at the start of draft decode for certain
+backends (e.g. ``ROCM_AITER_FA`` with eagle/eagle3 spec decode):
+
+    AssertionError: assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
+"""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from vllm.v1.worker.gpu.spec_decode import speculator as base_speculator
+from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator
+
+
+def _make_fake_speculator(
+    *,
+    max_num_reqs: int = 8,
+    max_num_tokens: int = 16,
+    max_model_len: int = 1024,
+    draft_max_seq_len: int = 1024,
+) -> SimpleNamespace:
+    """Build a fake EagleSpeculator with just the attributes used by
+    ``_build_draft_attn_metadata``. We deliberately avoid constructing a
+    real ``EagleSpeculator`` because that requires a full ``VllmConfig``
+    and a draft model.
+    """
+    fake_input_buffers = SimpleNamespace(
+        query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32),
+        seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
+    )
+    fake_block_tables = SimpleNamespace(
+        input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)],
+        slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64),
+    )
+    return SimpleNamespace(
+        arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),
+        block_tables=fake_block_tables,
+        input_buffers=fake_input_buffers,
+        attn_groups=[],
+        kv_cache_config=SimpleNamespace(kv_cache_groups=[]),
+        max_model_len=max_model_len,
+        draft_max_seq_len=draft_max_seq_len,
+    )
+
+
+def _run_build(fake, *, num_reqs, num_reqs_padded, num_tokens_padded, base, step):
+    captured: dict[str, object] = {}
+
+    def fake_build_attn_metadata(**kwargs):
+        captured.update(kwargs)
+        return {}
+
+    with patch.object(base_speculator, "build_attn_metadata", fake_build_attn_metadata):
+        EagleSpeculator._build_draft_attn_metadata(
+            fake,  # type: ignore[arg-type]
+            num_reqs=num_reqs,
+            num_reqs_padded=num_reqs_padded,
+            num_tokens_padded=num_tokens_padded,
+            seq_lens_cpu_upper_bound=base,
+            step=step,
+        )
+    return captured
+
+
+def test_build_draft_attn_metadata_sets_seq_lens_cpu_upper_bound():
+    """The fix: every per-step ``CommonAttentionMetadata`` carries a non-None
+    ``seq_lens_cpu_upper_bound`` derived from the target-side upper bound plus
+    the current draft-step offset. Padded entries are zeroed (matching the
+    main model runner's convention)."""
+    fake = _make_fake_speculator()
+    base = torch.tensor([100, 200, 300, 0], dtype=torch.int32)
+
+    captured = _run_build(
+        fake, num_reqs=3, num_reqs_padded=4, num_tokens_padded=4, base=base, step=2
+    )
+
+    bound = captured["seq_lens_cpu_upper_bound"]
+    assert isinstance(bound, torch.Tensor), (
+        "seq_lens_cpu_upper_bound must be a tensor, not None"
+    )
+    assert bound.shape == (4,), (
+        f"expected shape (num_reqs_padded=4,), got {bound.shape}"
+    )
+    assert bound.device.type == "cpu"
+    assert bound.dtype == torch.int32
+    # base[:num_reqs] + step, padded tail zeroed.
+    assert torch.equal(bound, torch.tensor([102, 202, 302, 0], dtype=torch.int32))
+
+
+def test_build_draft_attn_metadata_handles_zero_unpadded_reqs():
+    """Edge case: when ``num_reqs == 0`` the upper-bound tensor must
+    still be a valid all-zero tensor of length ``num_reqs_padded``."""
+    fake = _make_fake_speculator()
+    base = torch.zeros(2, dtype=torch.int32)
+
+    captured = _run_build(
+        fake, num_reqs=0, num_reqs_padded=2, num_tokens_padded=2, base=base, step=1
+    )
+
+    bound = captured["seq_lens_cpu_upper_bound"]
+    assert isinstance(bound, torch.Tensor)
+    assert bound.shape == (2,)
+    assert torch.equal(bound, torch.zeros(2, dtype=torch.int32))
+
+
+def test_build_draft_attn_metadata_clamps_to_max_model_len():
+    """The per-request upper bound (target bound + step) is clamped to the
+    model length so it never exceeds the allocated KV range."""
+    fake = _make_fake_speculator(max_model_len=1024)
+    base = torch.tensor([1023, 500], dtype=torch.int32)
+
+    captured = _run_build(
+        fake, num_reqs=2, num_reqs_padded=2, num_tokens_padded=2, base=base, step=3
+    )
+
+    bound = captured["seq_lens_cpu_upper_bound"]
+    # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.
+    assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
@@ -9,6 +9,7 @@ session is active. These tests verify that delegation and the session guard.

 import pytest

+from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.v1.worker.gpu_worker import Worker


@@ -21,29 +22,55 @@ class _RecordingEngine:
        self.finished = False
        self.reset_count = 0
        self.update_calls: list[dict] = []
+        self.seen_configs: list[VllmConfig] = []
+
+    def _record_config(self) -> None:
+        self.seen_configs.append(get_current_vllm_config())

    def start_weight_update(self) -> None:
+        self._record_config()
        self.started = True

    def update_weights(self, update_info: dict) -> None:
+        self._record_config()
        self.update_calls.append(update_info)
        if self.raise_on_update:
            raise ValueError("boom")

    def finish_weight_update(self) -> None:
+        self._record_config()
        self.finished = True

    def reset_weight_update_target(self) -> None:
        self.reset_count += 1


+class _RecordingModelRunner:
+    def __init__(self) -> None:
+        self.seen_config: VllmConfig | None = None
+
+    def reload_weights(self) -> None:
+        self.seen_config = get_current_vllm_config()
+
+
 def _make_worker(engine: _RecordingEngine | None) -> Worker:
    worker = object.__new__(Worker)
+    worker.vllm_config = VllmConfig()
    worker.weight_transfer_engine = engine
    worker._weight_update_active = False
    return worker


+def test_reload_weights_sets_current_config():
+    worker = _make_worker(None)
+    model_runner = _RecordingModelRunner()
+    worker.model_runner = model_runner  # type: ignore[assignment]
+
+    Worker.reload_weights(worker)
+
+    assert model_runner.seen_config is worker.vllm_config
+
+
 def test_start_update_finish_delegates_to_engine():
    engine = _RecordingEngine()
    worker = _make_worker(engine)
@@ -60,6 +87,7 @@ def test_start_update_finish_delegates_to_engine():
    assert engine.finished is True
    assert engine.reset_count == 1
    assert worker._weight_update_active is False
+    assert engine.seen_configs == [worker.vllm_config] * 3


 def test_double_start_raises():
@@ -1083,7 +1083,8 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
        return {}

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

+        # Check flash_attn_supports_kv_cache_dtype for fp8 support per FA version.
+        # Accept both equality checks and membership checks such as `in (3, 4)`.
+        if node.name == "flash_attn_supports_kv_cache_dtype":
+            for n in ast.walk(node):
+                if not (
+                    isinstance(n, ast.Compare)
+                    and len(n.ops) == 1
+                    and len(n.comparators) == 1
+                ):
+                    continue
+                is_version_compare = (
+                    isinstance(n.left, ast.Name) and n.left.id == "fa_version"
+                ) or (
+                    isinstance(n.left, ast.Call)
+                    and isinstance(n.left.func, ast.Name)
+                    and n.left.func.id == "get_flash_attn_version"
+                )
+                if not is_version_compare:
+                    continue
+
+                versions: list[Any] = []
+                comparator = n.comparators[0]
+                if isinstance(n.ops[0], ast.Eq) and isinstance(
+                    comparator, ast.Constant
+                ):
+                    versions = [comparator.value]
+                elif isinstance(n.ops[0], ast.In) and isinstance(
+                    comparator, (ast.Tuple, ast.List, ast.Set)
+                ):
+                    versions = [
+                        elt.value
+                        for elt in comparator.elts
+                        if isinstance(elt, ast.Constant)
+                    ]
+
+                fa3_supports_fp8 |= 3 in versions
+                fa4_supports_fp8 |= 4 in versions
+
        # Check flash_attn_supports_sinks - looks for `fa_version == 3/4`
        # or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`)
        if node.name == "flash_attn_supports_sinks":
@@ -1199,7 +1238,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
        },
        "fa4": {
            "compute_capability": fa4_compute_cap,
-            "supports_fp8": False,
+            "supports_fp8": fa4_supports_fp8,
            "supports_sink": fa4_supports_sinks,
        },
    }
@@ -1255,6 +1294,22 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:
 # ---------------------------------------------------------------------------


+def _apply_fp8_support(kv_cache_dtypes: str, supports_fp8: bool) -> str:
+    """Add or remove fp8 dtypes from a comma-separated kv_cache_dtypes string.
+
+    The base FLASH_ATTN backend lists fp8 dtypes in ``supported_kv_cache_dtypes``,
+    but actual support varies by FA version (e.g. FA2 has no fp8 support), so each
+    variant's row is normalized to match its own capability.
+    """
+    fp8_dtypes = {"fp8", "fp8_e4m3", "fp8_e5m2"}
+    dtypes = kv_cache_dtypes.split(", ")
+    base_dtypes = [dtype for dtype in dtypes if dtype not in fp8_dtypes]
+    if supports_fp8:
+        advertised_fp8_dtypes = [dtype for dtype in dtypes if dtype in fp8_dtypes]
+        return ", ".join(base_dtypes + advertised_fp8_dtypes)
+    return ", ".join(base_dtypes)
+
+
 def _expand_flash_attn_variants(
    all_backends: list[dict[str, Any]],
    fa_features: dict[str, dict[str, Any]],
@@ -1275,6 +1330,9 @@ def _expand_flash_attn_variants(
        fa2["_sort_key"] = "FLASH_ATTN"
        fa2["_sort_order"] = 0
        fa2["supports_sink"] = fa_features["fa2"]["supports_sink"]
+        fa2["kv_cache_dtypes"] = _apply_fp8_support(
+            backend["kv_cache_dtypes"], fa_features["fa2"]["supports_fp8"]
+        )

        # Create FA3 entry (uses parsed compute_capability from fa_utils)
        fa3 = backend.copy()
@@ -1284,11 +1342,9 @@ def _expand_flash_attn_variants(
        if fa_features["fa3"]["compute_capability"]:
            fa3["compute_capability"] = fa_features["fa3"]["compute_capability"]
        fa3["supports_sink"] = fa_features["fa3"]["supports_sink"]
-        if fa_features["fa3"]["supports_fp8"]:
-            base_dtypes = backend["kv_cache_dtypes"].split(", ")
-            fp8_dtypes = ["fp8", "fp8_e4m3", "fp8_e5m2"]
-            new_dtypes = [d for d in fp8_dtypes if d not in base_dtypes]
-            fa3["kv_cache_dtypes"] = ", ".join(base_dtypes + new_dtypes)
+        fa3["kv_cache_dtypes"] = _apply_fp8_support(
+            backend["kv_cache_dtypes"], fa_features["fa3"]["supports_fp8"]
+        )

        expanded.append(fa2)
        expanded.append(fa3)
@@ -1302,6 +1358,9 @@ def _expand_flash_attn_variants(
            if fa_features["fa4"].get("compute_capability"):
                fa4["compute_capability"] = fa_features["fa4"]["compute_capability"]
            fa4["supports_sink"] = fa_features["fa4"]["supports_sink"]
+            fa4["kv_cache_dtypes"] = _apply_fp8_support(
+                backend["kv_cache_dtypes"], fa_features["fa4"]["supports_fp8"]
+            )
            expanded.append(fa4)

    return expanded
@@ -835,12 +835,16 @@ def cutlass_scaled_mm_azp(


 def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
-    if cuda_device_capability < 90 or cuda_device_capability >= 110:
+    # CUTLASS grouped FP8 MoE uses the same sm100-named implementation for
+    # SM10x and SM11x. Thor reports SM101 with CUDA 12 and SM110 with CUDA 13,
+    # so keep this Python guard aligned with the C++ support query/dispatch.
+    if cuda_device_capability < 90 or cuda_device_capability >= 120:
        return False
    try:
        return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
-    except AttributeError:
-        # Return False on non-CUDA platforms where it is not available
+    except (AttributeError, RuntimeError):
+        # Return False on non-CUDA platforms where it is not available or not
+        # implemented for the current build.
        return False


@@ -2070,6 +2074,93 @@ def selective_scan_fwd(
    )


+def causal_conv1d_update_cpu_vec(
+    x: torch.Tensor,
+    conv_state: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None = None,
+    activation: str | None = None,
+    conv_state_indices: torch.Tensor | None = None,
+    query_start_loc: torch.Tensor | None = None,
+    pad_slot_id: int = 0,
+) -> torch.Tensor:
+    return torch.ops._C.causal_conv1d_update_cpu_vec(
+        x,
+        conv_state,
+        weight,
+        bias,
+        activation,
+        conv_state_indices,
+        query_start_loc,
+        pad_slot_id,
+    )
+
+
+def selective_state_update_cpu(
+    state: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None,
+    z: torch.Tensor | None,
+    dt_bias: torch.Tensor | None,
+    dt_softplus: bool,
+    state_batch_indices: torch.Tensor | None,
+    dst_state_batch_indices: torch.Tensor | None,
+    null_block_id: int,
+    out: torch.Tensor,
+    num_accepted_tokens: torch.Tensor | None,
+    cu_seqlens: torch.Tensor | None,
+):
+    torch.ops._C.selective_state_update_cpu(
+        state,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D,
+        z,
+        dt_bias,
+        dt_softplus,
+        state_batch_indices,
+        dst_state_batch_indices,
+        null_block_id,
+        out,
+        num_accepted_tokens,
+        cu_seqlens,
+    )
+
+
+def mamba_chunk_scan_fwd_cpu(
+    out: torch.Tensor,
+    final_states: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    D: torch.Tensor | None,
+    z: torch.Tensor | None,
+    cu_seqlens: torch.Tensor,
+) -> None:
+    """Prefill SSM scan kernel. out and final_states are written in-place."""
+    torch.ops._C.mamba_chunk_scan_fwd_cpu(
+        out,
+        final_states,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        D,
+        z,
+        cu_seqlens,
+    )
+
+
 # ROCm skinny gemms
 def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
    return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
@@ -3776,10 +3867,10 @@ def fusedQuantizeMx(
        raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'")


-if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"):
+if hasattr(torch.ops._qutlass_C, "fusedQuantizeNvAbsMax"):

-    @register_fake("_qutlass_C::fusedQuantizeNv")
-    def _fake_fused_quantize_nv(
+    @register_fake("_qutlass_C::fusedQuantizeNvAbsMax")
+    def _fake_fused_quantize_nv_absmax(
        a: torch.Tensor,
        b: torch.Tensor,
        xh_e2m1: torch.Tensor,
@@ -3805,7 +3896,9 @@ def fusedQuantizeNv(
        padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device
    )

-    return torch.ops._qutlass_C.fusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale)
+    return torch.ops._qutlass_C.fusedQuantizeNvAbsMax(
+        a, b, xh_e2m1, xh_e4m3, global_scale
+    )


 def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
@@ -219,6 +219,63 @@ def _xpu_ops_deepseek_scaling_rope_fake(
    return query, key


+def _xpu_fp8_bmm_impl(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    out_dtype: torch.dtype,
+    a_scale: torch.Tensor,
+    b_scale: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> torch.Tensor:
+    """XPU FP8 batched GEMM implementation for ``torch.ops.vllm.xpu_fp8_bmm``.
+
+    Computes batched matrix multiplication over the leading group dimension:
+    ``[G, M, K] @ [G, K, N] -> [G, M, N]``.
+
+    Args:
+        a: FP8 activation tensor with shape ``[G, M, K]``.
+            Does not need to be contiguous.
+        b: FP8 weight tensor with shape ``[G, K, N]``.
+            Does not need to be contiguous.
+        out_dtype: Output dtype accepted by the kernel (typically
+            ``torch.bfloat16`` for the DeepSeek-V4 O-proj path).
+        a_scale: Activation scale tensor for ``a``.
+            In current DeepSeek-V4 XPU usage it is block-scaled with shape
+            ``[G, M, K/bs]`` (``bs`` is the quant block size, e.g. 128).
+            Must be contiguous.
+        b_scale: Weight scale tensor for ``b``.
+            In current DeepSeek-V4 XPU usage it is block-scaled with shape
+            ``[G, K/bs, N/bs]`` (``bs`` is the quant block size, e.g. 128).
+            Must be contiguous.
+        bias: Optional bias tensor. Pass ``None`` when no bias is required.
+
+    Returns:
+        Output tensor with shape ``[G, M, N]`` and dtype ``out_dtype``.
+
+    Notes:
+        This implementation centralizes access to
+        ``torch.ops._xpu_C.fp8_bmm``. Both scales must be contiguous, while
+        ``a`` and ``b`` may be non-contiguous views.
+    """
+    return torch.ops._xpu_C.fp8_bmm(a, b, out_dtype, a_scale, b_scale, bias)
+
+
+def _xpu_fp8_bmm_fake(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    out_dtype: torch.dtype,
+    a_scale: torch.Tensor,
+    b_scale: torch.Tensor,
+    bias: torch.Tensor | None,
+) -> torch.Tensor:
+    # [G, M, K] @ [G, K, N] => [G, M, N]
+    return torch.empty(
+        (a.shape[0], a.shape[1], b.shape[2]),
+        dtype=out_dtype,
+        device=a.device,
+    )
+
+
 def _xpu_fp8_mqa_logits_impl(
    q: torch.Tensor,
    k_quant: torch.Tensor,
@@ -1053,6 +1110,12 @@ class xpu_ops:
                fake_impl=_xpu_mxfp4_quantize_fake,
            )

+            direct_register_custom_op(
+                op_name="xpu_fp8_bmm",
+                op_func=_xpu_fp8_bmm_impl,
+                fake_impl=_xpu_fp8_bmm_fake,
+            )
+
            direct_register_custom_op(
                op_name="xpu_fp8_mqa_logits",
                op_func=_xpu_fp8_mqa_logits_impl,
@@ -175,9 +175,15 @@ class KernelConfig:
    enable_flashinfer_autotune: bool = None  # type: ignore[assignment]
    """If True, run FlashInfer autotuning during kernel warmup."""

+    # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
+    # to the shared JIT warmup infrastructure.
+    # https://github.com/vllm-project/vllm/pull/47451
    enable_cutedsl_warmup: bool = True
    """If True, run CuTeDSL compile warmup during kernel warmup."""

+    enable_jit_warmup: bool = True
+    """If True, run JIT compile warmup during kernel warmup."""
+
    enable_bf16x3_router_gemm: bool = False
    """If True, use the experimental SM100 BF16x3 CuteDSL router GEMM."""

@@ -250,6 +256,7 @@ class KernelConfig:
        """
        ignored_factors = {
            "enable_cutedsl_warmup",
+            "enable_jit_warmup",
            "enable_flashinfer_autotune",
            "ir_op_priority",  # handled separately below
        }
@@ -260,6 +267,7 @@ class KernelConfig:
    @field_validator(
        "enable_flashinfer_autotune",
        "enable_cutedsl_warmup",
+        "enable_jit_warmup",
        mode="wrap",
    )
    @classmethod
@@ -27,6 +27,7 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta):

    TRITON = "triton"
    FLASHINFER = "flashinfer"
+    CPU = "cpu"


@config
@@ -1935,13 +1935,8 @@ class NixlBaseConnectorWorker:
        indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)

        for _, cache_or_caches in self.device_kv_caches.items():
-            blocks_to_update = cache_or_caches.index_select(1, indices)
            current_platform.pack_kv_cache(
-                key=blocks_to_update[0],
-                value=blocks_to_update[1],
-                key_cache=cache_or_caches[0],
-                value_cache=cache_or_caches[1],
-                block_ids=block_ids,
+                kv_cache=cache_or_caches,
                indices=indices,
            )

@@ -61,9 +61,9 @@ def get_offloading_event_group_spec(

@dataclass(slots=True)
 class _OffloadEventMetadata:
-    """BlockStored payload snapshot for one OffloadKey, captured at store
-    time and kept until the matching eviction event. ``medium`` is forwarded
-    from the OffloadingEvent."""
+    """BlockStored payload snapshot for one OffloadKey, captured while the
+    Request is available and kept until the matching eviction event. ``medium``
+    is forwarded from the OffloadingEvent."""

    # The chunk's constituent block hashes; the last one is the OffloadKey.
    block_hashes: tuple[BlockHash, ...]
@@ -81,10 +81,11 @@ class _OffloadEventMetadata:
 class OffloadingEventsTracker:
    """Tracks offloaded chunks' KV event payloads from store to eviction.

-    The scheduler calls :meth:`record_store` from ``_build_store_jobs``
-    while the ``Request`` is available, and routes the manager's raw
-    :class:`OffloadingEvent` stream through :meth:`take_events`. All state
-    is bounded by the CPU pool capacity and cleared by :meth:`reset`.
+    The scheduler calls :meth:`record_store` from ``_build_store_jobs`` and
+    :meth:`record_lookup` for ready primary-tier hits while the ``Request`` is
+    available. Deferred and missing lookups add no state. Under the connector's
+    supported success-only transfer model, entries follow primary allocations
+    until CPU removal translation or :meth:`reset`.
    """

    def __init__(self, config: OffloadingKVEventsConfig):
@@ -93,8 +94,7 @@ class OffloadingEventsTracker:
            config.enable_kv_cache_events and config.self_describing_kv_events
        )

-        # OffloadKey -> payload snapshot, kept until the eviction event so
-        # BlockRemoved can fan out. Bounded: one entry per offloaded chunk.
+        # OffloadKey -> payload snapshot, kept until CPU removal or reset.
        self._pending_event_metadata: dict[OffloadKey, _OffloadEventMetadata] = {}

    def record_store(
@@ -116,6 +116,23 @@ class OffloadingEventsTracker:
        meta = self._build_event_metadata(req, group_config, chunk_idx)
        self._pending_event_metadata[offload_key] = meta

+    def record_lookup(
+        self,
+        req: Request,
+        group_config: "GroupOffloadConfig",
+        chunk_idx: int,
+        offload_key: OffloadKey,
+    ) -> None:
+        """Snapshot metadata for a ready primary-tier lookup hit."""
+        if not self.self_describing_enabled:
+            return
+        if group_config.sliding_window_size_in_chunks is not None:
+            return
+        if offload_key not in self._pending_event_metadata:
+            self._pending_event_metadata[offload_key] = self._build_event_metadata(
+                req, group_config, chunk_idx
+            )
+
    def take_events(self, events: Iterable[OffloadingEvent]) -> Iterable[KVCacheEvent]:
        """Translate raw OffloadingEvents into self-describing KV events.

@@ -165,7 +182,7 @@ class OffloadingEventsTracker:
        assert len(chunk_hashes) == hbf

        if group_config.sliding_window_size_in_chunks is not None:
-            # record_store filters these out before calling this helper.
+            # The recording methods filter these out before calling this helper.
            raise AssertionError("self-describing events only support full attention")

        parent_block_hash: BlockHash | None
@@ -232,7 +249,8 @@ class OffloadingEventsTracker:
                        "OffloadingEventsTracker: no event metadata for "
                        "offload key during BlockStored emission; emitting a "
                        "placeholder payload. Expected for non-full-attention "
-                        "groups; otherwise indicates a missing populate path."
+                        "groups and promotions not observed as a primary-tier "
+                        "hit before translation."
                    )
                yield self._placeholder_stored(key, event.medium, locality)
                continue
@@ -111,6 +111,26 @@ def get_sliding_window_size_in_chunks(
    return None


+def is_store_reachable_swa_chunk(
+    absolute_chunk_index: int,
+    storable_chunk_count: int,
+    alignment_chunk_count: int | None,
+    sliding_window_chunks: int | None,
+    is_eagle_group: bool,
+) -> bool:
+    """Return whether an SWA chunk can participate in an external-cache hit."""
+    if alignment_chunk_count is None:
+        return True
+    assert sliding_window_chunks is not None
+    position_in_segment = absolute_chunk_index % alignment_chunk_count
+    segment_start = absolute_chunk_index - position_in_segment
+    actual_segment_length = min(
+        alignment_chunk_count, storable_chunk_count - segment_start
+    )
+    reachable_tail = sliding_window_chunks + int(is_eagle_group)
+    return position_in_segment >= actual_segment_length - reachable_tail
+
+
 def resolve_mamba_align_size(
    spec: "OffloadingSpec", kv_cache_config: KVCacheConfig
 ) -> int | None:
@@ -463,15 +483,27 @@ class OffloadingConnectorScheduler:
            del self._req_status[req_id]

    def _maximal_prefix_lookup(
-        self, keys: Iterable[OffloadKey], req_context: ReqContext
+        self,
+        keys: Iterable[OffloadKey],
+        req_context: ReqContext,
+        req: Request,
+        group_config: GroupOffloadConfig,
+        start_chunk_idx: int,
    ) -> int | None:
        """Return the number of consecutive offloaded chunks from the start,
        or None if the backend deferred a lookup."""
        hit_count = 0
        defer_lookup = False
-        for key in keys:
-            match self.manager.lookup(key, req_context):
+        for local_idx, key in enumerate(keys):
+            result = self.manager.lookup(key, req_context)
+            match result:
                case LookupResult.HIT:
+                    self._events_tracker.record_lookup(
+                        req,
+                        group_config,
+                        start_chunk_idx + local_idx,
+                        key,
+                    )
                    hit_count += 1
                case LookupResult.HIT_PENDING:
                    defer_lookup = True
@@ -616,7 +648,11 @@ class OffloadingConnectorScheduler:
                num_hit_chunks: int | None
                if sliding_window_size_in_chunks is None:
                    num_hit_chunks = self._maximal_prefix_lookup(
-                        offload_keys, req_status.req_context
+                        offload_keys,
+                        req_status.req_context,
+                        req_status.req,
+                        group_config,
+                        start_chunk_idx,
                    )
                else:
                    required_window = sliding_window_size_in_chunks
@@ -974,9 +1010,6 @@ class OffloadingConnectorScheduler:
                ]
                assert len(offload_keys) == len(offload_block_ids)

-                alignment_chunk_count = group_config.alignment_chunk_count
-                tail = group_config.sliding_window_size_in_chunks
-
                for key_idx, (offload_key, block_id) in enumerate(
                    zip(offload_keys, offload_block_ids)
                ):
@@ -984,15 +1017,18 @@ class OffloadingConnectorScheduler:
                        continue
                    # Skip SWA chunks that can never serve a load hit:
                    # within each full-attention alignment segment, only the
-                    # trailing `tail` chunks are reachable by
-                    # _sliding_window_lookup. For DeepSeek V4 with 100K
-                    # tokens this reduces SWA stores by ~78%.
-                    if alignment_chunk_count is not None:
-                        assert tail is not None
-                        abs_chunk_idx = start_chunk_idx + key_idx
-                        pos_in_segment = abs_chunk_idx % alignment_chunk_count
-                        if pos_in_segment < alignment_chunk_count - tail:
-                            continue
+                    # trailing chunks queried by _sliding_window_lookup are
+                    # reachable. EAGLE/MTP requires one additional chunk that
+                    # lookup later drops as its volatile draft tail.
+                    abs_chunk_idx = start_chunk_idx + key_idx
+                    if not is_store_reachable_swa_chunk(
+                        abs_chunk_idx,
+                        num_chunks,
+                        group_config.alignment_chunk_count,
+                        group_config.sliding_window_size_in_chunks,
+                        group_config.is_eagle_group,
+                    ):
+                        continue
                    new_offload_keys.append(offload_key)

            if not new_offload_keys:
@@ -400,13 +400,10 @@ class GroupCoordinator:
        self.rank = torch.distributed.get_rank()
        self.local_rank = local_rank
        self.device_index: int
-        if _WORLD is not None:
-            self.device_index = _WORLD.device_index
-        else:
-            assert local_rank >= 0, (
-                "local_rank must be provided when creating the world group"
-            )
-            self.device_index = local_rank
+        assert local_rank >= 0, (
+            "local_rank must be provided when creating the world group"
+        )
+        self.device_index = local_rank

        self_device_group = None
        self_cpu_group = None
@@ -1,11 +1,68 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
+import os
+import sys
+from pathlib import Path

-from vllm.benchmarks.serve import add_cli_args, main
+from vllm.benchmarks.serve import add_cli_args
+from vllm.benchmarks.serve import main as python_main
 from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
+from vllm.logger import init_logger
 from vllm.utils.argparse_utils import FlexibleArgumentParser

+logger = init_logger(__name__)
+_RUST_CLI_PATH = Path(__file__).resolve().parents[3] / "vllm-rs"
+_RUST_SUPPORTED_DATASETS = frozenset(
+    {
+        "custom",
+        "hf",
+        "prefix_repetition",
+        "random",
+        "random-mm",
+        "random-rerank",
+        "sharegpt",
+        "sonnet",
+        "speed_bench",
+    }
+)
+_RUST_SUPPORTED_BACKENDS = frozenset(
+    {
+        "openai",
+        "openai-chat",
+        "openai-embeddings",
+        "openai-embeddings-chat",
+        "vllm",
+        "vllm-pooling",
+        "vllm-rerank",
+    }
+)
+
+
+def _rust_unsupported_reason(args: argparse.Namespace) -> str | None:
+    if args.dataset_name not in _RUST_SUPPORTED_DATASETS:
+        return f"dataset {args.dataset_name!r} is not supported by the Rust benchmark"
+    if args.backend not in _RUST_SUPPORTED_BACKENDS:
+        return f"backend {args.backend!r} is not supported by the Rust benchmark"
+    return None
+
+
+def _maybe_exec_rust_bench(args: argparse.Namespace) -> None:
+    if reason := _rust_unsupported_reason(args):
+        logger.info("Using Python benchmark: %s.", reason)
+        return
+
+    if not _RUST_CLI_PATH.is_file():
+        logger.warning(
+            "Rust benchmark binary not found at %s; falling back to Python.",
+            _RUST_CLI_PATH,
+        )
+        return
+
+    rust_cli = str(_RUST_CLI_PATH)
+    logger.info("Delegating `vllm bench serve` to Rust binary at %s.", rust_cli)
+    os.execv(rust_cli, [rust_cli, "bench", "serve", *sys.argv[3:]])
+

 class BenchmarkServingSubcommand(BenchmarkSubcommandBase):
    """The `serve` subcommand for `vllm bench`."""
@@ -19,4 +76,5 @@ class BenchmarkServingSubcommand(BenchmarkSubcommandBase):

    @staticmethod
    def cmd(args: argparse.Namespace) -> None:
-        main(args)
+        _maybe_exec_rust_bench(args)
+        python_main(args)
@@ -757,6 +757,18 @@ class ChatCompletionRequest(OpenAIBaseModel):
                parameter="logprob_token_ids",
            )

+        # These fields are integers, but `mode="before"` runs on the raw
+        # request data, so a non-numeric value (e.g. a JSON string) would
+        # reach the comparisons below and raise TypeError -> HTTP 500. Reject
+        # it here so the client gets a clean 400 instead.
+        for field_name in ("prompt_logprobs", "top_logprobs"):
+            field_value = data.get(field_name)
+            if field_value is not None and not isinstance(field_value, (int, float)):
+                raise VLLMValidationError(
+                    f"`{field_name}` must be an integer.",
+                    parameter=field_name,
+                    value=field_value,
+                )
        if (prompt_logprobs := data.get("prompt_logprobs")) is not None:
            if data.get("stream") and (prompt_logprobs > 0 or prompt_logprobs == -1):
                raise VLLMValidationError(
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Shengqi ChenandGitHub	a02155c787	Merge branch 'main' into cuda-arch-fixup	2026-07-21 09:33:48 +08:00
Chris LeonardandGitHub	97a98006b0	Update qutlass cmake for stable abi (#47879 ) Signed-off-by: Chris Leonard <chleonar@redhat.com>	2026-07-20 18:31:10 -07:00
Alejandro Paredes La Torre GitHub Shengqi Chen	0a684ab0c0	[Bugfix] Fix WSL circular import from pin_memory warning_once (#48444 ) Signed-off-by: AlejandroParedesLT <alejandroparedeslatorre@gmail.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>	2026-07-20 18:30:55 -07:00
flutist-alibaba GitHub Copilot Autofix powered by AI Shengqi Chen	0d9210a502	Fixes non-coalesced HBM access in marlin_int4_fp8_preprocess_kernel_awq (#47268 ) Signed-off-by: xjx <493337577@qq.com> Signed-off-by: flutist-alibaba <30485581+flutist@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>	2026-07-20 18:30:38 -07:00
Lena Onyshchenko GitHub Cursor	1d874867ea	[Misc][Docs] Fix broken protocol link in speech_to_text doc (#47212 ) Signed-off-by: oonyshch <xonyshch@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>	2026-07-21 01:05:05 +00:00
ormandj GitHub Luke Alonso Martin Vit Claude Yifan Qiao	2e2e626b40	[Bugfix] Count per-group blocks in get_max_concurrency_for_kv_cache_config (#48317 ) Signed-off-by: David Orman <ormandj@corenode.com> Co-authored-by: Luke Alonso <lalonso@gmail.com> Co-authored-by: Martin Vit <martin@voipmonitor.org> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>	2026-07-21 00:29:10 +00:00
Nick HillandGitHub	af91f4b3e4	[Cleanup] Remove unused StructuredOutputRequest.status field (#49235 ) Signed-off-by: Nick Hill <nickhill123@gmail.com>	2026-07-21 00:06:16 +00:00
sungsoo ha GitHub Claude Opus 4.8 Lucas Wilkinson Matthew Bonanni	2396a61108	[Attention][MLA][DCP] Query replication for MLA decode (DeepSeek-V2/R1 + Kimi-K2.5) (#45964 ) Signed-off-by: Sungsoo Ha <sungsooh@nvidia.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>	2026-07-20 23:51:27 +00:00
Xuanyu Zhang GitHub Claude aoshen02 OpenAI Codex	97a668152b	[RL Infra][FlashInfer] Enable router replay output from FlashInfer monolithic MoE kernel (#44214 ) Signed-off-by: Xuanyu Zhang <xuanyu.zhang@mistral.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com>	2026-07-20 16:45:10 -07:00
Matthew Bonanni GitHub haosdent Claude Opus 4.7 Justin Wood	58b2012aa2	[copy of #45208 ] CuMem slept-L1 fragmentation accounting (#49208 ) Signed-off-by: haosdent <haosdent@gmail.com> Signed-off-by: Justin Wood <justin.m.wood@me.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: haosdent <haosdent@gmail.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Justin Wood <jwood@me.com>	2026-07-20 23:08:11 +00:00
Ning XieandGitHub	b7c20d0cfa	[chore] adjust logo be more friendly to white background terminal (#48938 ) Signed-off-by: Andy Xie <andy.xning@gmail.com>	2026-07-20 15:15:28 -07:00
TJianandGitHub	a2b1f9fc3b	[ROCm] [Release] [Bugfix] Fix the per commit wheel release pipeline. (#49245 ) Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>	2026-07-20 22:12:38 +00:00
Fynn Schmitt-Ulms GitHub Michael Goin	642076d26c	Support loading sample_from_anchor flag from speculators config (#48639 ) Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>	2026-07-20 14:52:43 -07:00
Charlie FuandGitHub	5feb3950e5	[ROCm][CI] fix test_rocm_quick_reduce.py (#49234 ) Signed-off-by: charlifu <charlifu@amd.com>	2026-07-20 16:39:33 -05:00
Oxana Korzh GitHub Cursor okorzh-amd Matthew Bonanni	4ec199b66a	[Bugfix][Spec-Decode] Populate draft seq_lens_cpu_upper_bound for spec-decode attention metadata (#44492 ) Signed-off-by: Oxana Korzh <okorzh@amd.com> Signed-off-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: okorzh-amd <okorzh-amd@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>	2026-07-20 20:50:58 +00:00
Roberto L. Castro GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	7ca017778f	[Feat][Perf] Add new warmup infrastructure for JITs (#47451 ) Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-20 13:21:55 -07:00
coltonottley GitHub Colton Ottley Or Ozeri	fbfe58133d	[Bugfix][KV Offload] Preserve reachable tails for hybrid SWA groups (#48911 ) Signed-off-by: Colton Ottley <colton@ottleyengineering.com> Co-authored-by: Colton Ottley <colton@ottleyengineering.com> Co-authored-by: Or Ozeri <oro@il.ibm.com>	2026-07-20 22:12:36 +03:00
wkutak GitHub Isotr0py Roger Wang	9dd62d80ab	Cosmos3 FP8 ModelOpt/Diffusers remapping (#48952 ) Signed-off-by: Wojciech Kutak <wkutak@nvidia.com> Signed-off-by: wkutak <wkutak@nvidia.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Roger Wang <hey@rogerw.io>	2026-07-20 11:31:48 -07:00
tc-mb GitHub wjinxu	f878367898	[Revert][Bugfix] Restore MiniCPM-V 4.6 ViT QKV weight loader (#49193 ) Signed-off-by: wjinxu <1299461899@qq.com> Co-authored-by: wjinxu <1299461899@qq.com>	2026-07-20 18:17:46 +00:00
Matthew Bonanni GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	bd091079cb	[Attention] FlashAttention 4 SM100 FP8 kv cache support (#42569 ) Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-20 10:53:27 -07:00
Liangqiusong GitHub Kunshang Ji	b23bd73f54	[XPU]add sycl path for Mhc (#47245 ) Signed-off-by: root <xiaolong.guo@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-07-20 15:32:54 +00:00
Bugen ZhaoandGitHub	e2d7adeb64	[Rust Frontend] Bump `xgrammar-structural-tag` and enable local extension (#49161 ) Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-20 16:22:24 +01:00
Isotr0pyandGitHub	15cb8e140d	[Multimodal] Allow keeping original image mode for ImageIO (#49159 ) Signed-off-by: Isotr0py <Isotr0py@outlook.com>	2026-07-20 13:42:45 +00:00
Chang Guo GitHub OpenAI Codex	f007cceb42	[KV Offload] Support self-describing KV events with TieringOffloadingSpec (#48679 ) Signed-off-by: Change72 <changg@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-20 16:41:58 +03:00
wangqian GitHub OpenAI Codex	0a5069e4e3	[Bugfix][Gemma4] Fix ModelOpt mixed-precision MoE config mapping (#48563 ) Signed-off-by: wangqian <601731555@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-20 06:39:28 -07:00
EdalatiAli GitHub Cursor Nicolò Lucchesi	8ce53a616e	[Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches (#47574 ) Signed-off-by: EdalatiAli <aliedalati@cohere.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>	2026-07-20 13:18:17 +00:00
Lena OnyshchenkoandGitHub	ae10e855ab	[Misc][Docs] Remove duplicate CodeGeex4 row in XPU model table (#47210 ) Signed-off-by: oonyshch <xonyshch@gmail.com>	2026-07-20 10:05:36 +00:00
hclandGitHub	530ee36a0d	fix(openai): reject non-numeric logprobs with 400 instead of 500 (#49144 ) Signed-off-by: Chenglun Hu <chenglunhu@gmail.com>	2026-07-20 10:04:50 +00:00
Salt SatoandGitHub	d835ad572c	[Bugfix][Rust Frontend] Map missing prompt logprobs for single-token prompts in chat and raw generate (#49111 ) Signed-off-by: Feathbow <feathbow@gmail.com>	2026-07-20 10:00:06 +00:00
Lena Onyshchenko GitHub Cursor	47d0597ca2	[Misc][Docs] Fix broken csrc kernel links in fusions doc (#47211 ) Signed-off-by: oonyshch <xonyshch@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>	2026-07-20 09:44:25 +00:00
ReidandGitHub	818cf61e91	[Rust Frontend] Fix macro-based content format detection (#49042 ) Signed-off-by: reidliu41 <reid201711@gmail.com>	2026-07-20 09:39:13 +00:00
Bugen Zhao GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	c01618fdc8	[Rust][Benchmark] Integrate `vllm-bench` to `vllm-rs` & `vllm` CLI (#48930 ) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>	2026-07-20 09:31:25 +00:00
Xiaochang Wu GitHub Kunshang Jimergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	823eaf667d	[XPU] FP8 o_proj with fp8_bmm and load-time scale transpose (#48334 ) Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	2026-07-20 16:32:03 +08:00
Sage GitHubmergify[bot] <37929162+mergify[bot]@users.noreply.github.com>	f1f1259692	[Rust Frontend] Use zero-copy slicing for multimodal tensors (#48781 ) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Sage Ahrac <sagiahrak@gmail.com>	2026-07-20 16:28:25 +08:00
zofia GitHub mayuyuacemergify[bot] <37929162+mergify[bot]@users.noreply.github.com>Kunshang Ji	df13b5aef5	[XPU] [MoE] add quant input when prepare for fusedmoe (#47122 ) Signed-off-by: mayuyuace <qiming1.zhang@intel.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Signed-off-by: zofia <110436990+zufangzhu@users.noreply.github.com> Co-authored-by: mayuyuace <qiming1.zhang@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-07-20 15:47:26 +08:00
Sihan Chen GitHubLi, Jiang <jiang1.li@intel.com>	4938d44a3b	[CPU] fixes heterogeneous NIXL KV transfer into CPU_ATTN decode workers (#47871 ) Signed-off-by: Spycsh <sihan.chen@intel.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>	2026-07-20 07:33:13 +00:00
Michał Ganczarenko GitHub Kunshang Ji	37bf988c2f	[XPU][Bugfix] Fix GroupCoordinator device_index (#47295 ) Signed-off-by: Michal Ganczarenko <michal.ganczarenko@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>	2026-07-20 15:25:56 +08:00
aoshen02andGitHub	9459fc6471	[Bugfix][RL] Set vLLM config during weight reload (#45989 ) Signed-off-by: aoshen02 <aoshen@inferact.ai>	2026-07-20 15:02:56 +08:00
Itay Etelis GitHub Itay Etelis	5245c80564	[Doc] Document blocks_per_chunk in the KV offloading guide (#49100 ) Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com>	2026-07-20 09:48:43 +03:00
Yifan Qiao GitHub OpenAI Codex	9bc266d923	[Bugfix][KV Offload] Propagate EAGLE mode to SimpleCPU coordinator (#49071 ) Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>	2026-07-20 06:39:11 +00:00
Akash kaothalkar GitHub Akash Kaothalkar Akash Kaothalkar Akash Kaothalkar Akash kaothalkarLi, Jiang <jiang1.li@intel.com>	5c9f6557d7	[Hardware][CPU] Enable granite-4 model on cpu (#47641 ) Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Signed-off-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com> Signed-off-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan> Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com> Co-authored-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan> Co-authored-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com> Co-authored-by: Akash kaothalkar <akash.kaothalkar@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>	2026-07-20 06:15:16 +00:00
Shengqi ChenandGitHub	319db65b68	Merge branch 'main' into cuda-arch-fixup	2026-07-09 01:55:33 +08:00
Shengqi ChenandGitHub	83a7669827	Merge branch 'main' into cuda-arch-fixup Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-07-07 00:13:53 +08:00
Shengqi ChenandGitHub	401bed48ad	Merge branch 'main' into cuda-arch-fixup Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-07-03 15:54:56 +08:00
Shengqi ChenandGitHub	d492d1e697	Merge branch 'main' into cuda-arch-fixup	2026-06-30 21:30:29 +08:00
Shengqi Chen	f37e113590	[Build] Apply ruff format to CUDA arch regex Keep the compiled CUDA arch regex on one line to match ruff-format output. Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-06-30 20:37:59 +08:00
Shengqi Chen	a10e369f06	[Build] Address CUDA arch review comments Fix CUDA arch warning tests so they do not depend on importing the stable libtorch extension, which keeps the warning coverage active in lightweight CI environments and avoids mypy treating a fixture value as a base class. Use regex for the compiled-arch parser, apply formatter output, and make the CUTLASS grouped GEMM Python support query fall back to false when the op is unavailable or unimplemented in the current build. Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-06-30 20:03:28 +08:00
Shengqi Chen	aa3f2efe42	[Temp] Cherry-pick #47139 to fix build Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-06-30 19:38:28 +08:00
Shengqi ChenandCodex	a32d1bff95	[Doc] Document CUDA wheel architecture coverage Explain that pre-built CUDA wheels use the architecture lists selected by the release and build pipelines, which may be narrower than the full set vLLM can build from source. Call out CUDA 12.9 architecture-specific wheel coverage, CUDA 13 family-specific targets, and the no-kernel-image error users may see when a wheel does not cover their GPU. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-06-30 19:27:42 +08:00
Shengqi ChenandCodex	ee47a21fcd	[Build] Warn on uncovered CUDA device architectures Expose the compiled CUDA arch list from the stable extension and check visible CUDA devices against it during CUDA platform startup. The runtime check distinguishes exact architecture targets from CUDA 13 family targets so users get an early warning before hitting missing kernel images. The startup warning is limited to the NVML-backed CUDA platform path to preserve the existing no-CUDA-init import behavior for non-NVML environments. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-06-30 19:27:42 +08:00
Shengqi ChenandCodex	0f471a3088	[Build] Scope stable CUDA kernel feature macros Keep optional stable CUDA kernel feature macros on the source files that consume them instead of adding them to VLLM_GPU_FLAGS. This avoids perturbing unrelated compile commands and invalidating more cache entries when optional kernel families change. Also align CUTLASS grouped MoE support with the SM10x/SM11x family so Thor works under both CUDA 12 SM101 and CUDA 13 SM110 reporting, and remove the stale ENABLE_CUTLASS_MLA definition left after the old CUTLASS MLA path was deleted. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Shengqi Chen <harry-chen@outlook.com>	2026-06-30 19:27:42 +08:00