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
109 changed files with 4960 additions and 1111 deletions
@@ -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
@@ -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)
@@ -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
@@ -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 |
@@ -75,7 +75,7 @@ vllm serve <model> \
 | `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, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
+| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. 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
@@ -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:
@@ -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
@@ -6321,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",
@@ -143,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",
@@ -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);
@@ -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
    }

@@ -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)


@@ -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))
@@ -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


@@ -3863,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,
@@ -3892,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:
@@ -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
@@ -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:
@@ -325,13 +325,12 @@ def log_version_and_model(lgr: Logger, version: str, model_name: str) -> None:
        message = "vLLM server version %s, serving model %s"
    else:
        logo_template = Template(
-            "\n       ${w}█     █     █▄   ▄█${r}\n"
-            " ${o}▄▄${r} ${b}▄█${r} ${w}█     █     █ ▀▄▀ █${r}  version ${w}%s${r}\n"
-            "  ${o}█${r}${b}▄█▀${r} ${w}█     █     █     █${r}  model   ${w}%s${r}\n"
-            "   ${b}▀▀${r}  ${w}▀▀▀▀▀ ▀▀▀▀▀ ▀     ▀${r}\n"
+            "\n       ${b}█     █     █▄   ▄█${r}\n"
+            " ${o}▄▄${r} ${b}▄█${r} ${b}█     █     █ ▀▄▀ █${r}  version ${b}%s${r}\n"
+            "  ${o}█${r}${b}▄█▀${r} ${b}█     █     █     █${r}  model   ${b}%s${r}\n"
+            "   ${b}▀▀${r}  ${b}▀▀▀▀▀ ▀▀▀▀▀ ▀     ▀${r}\n"
        )
        colors = {
-            "w": "\033[97;1m",  # white
            "o": "\033[93m",  # orange
            "b": "\033[94m",  # blue
            "r": "\033[0m",  # reset
@@ -186,6 +186,7 @@ if TYPE_CHECKING:
    VLLM_MOE_USE_DEEP_GEMM: bool = True
    VLLM_USE_DEEP_GEMM_E8M0: bool = True
    VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True
+    VLLM_DCP_Q_REPLICATE: bool = False
    VLLM_DEEP_GEMM_WARMUP: Literal[
        "skip",
        "full",
@@ -1490,6 +1491,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES": lambda: bool(
        int(os.getenv("VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES", "1"))
    ),
+    # Opt-in MLA DCP query replication: skip the decode query all-gather.
+    "VLLM_DCP_Q_REPLICATE": lambda: bool(int(os.getenv("VLLM_DCP_Q_REPLICATE", "0"))),
    # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm
    # JIT all the required kernels before model execution so there is no
    # JIT'ing in the hot-path. However, this warmup increases the engine
@@ -366,6 +366,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        q_lora_rank: int | None,
        kv_lora_rank: int,
        kv_b_proj: ColumnParallelLinear,
+        dcp_q_replicate: bool = False,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
@@ -384,10 +385,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.kv_b_proj = kv_b_proj
+        self.dcp_q_replicate = dcp_q_replicate
+        self.W_UK_T_dcp_qrep: torch.Tensor | None = None
        self.head_size = kv_lora_rank + qk_rope_head_dim
        self.layer_name = prefix
        self.indexer = indexer
-
        self.num_kv_heads = 1
        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

@@ -591,6 +593,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        kv_c_normed: torch.Tensor,
        k_pe: torch.Tensor,
        output_shape: torch.Size | None = None,
+        q_dcp_replicated: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if self.calculate_kv_scales:
            torch.ops.vllm.maybe_calc_kv_scales(
@@ -646,6 +649,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                self_kv_cache,
                attn_metadata,
                output=output,
+                q_dcp_replicated=q_dcp_replicated,
            )
            return output
        else:
@@ -665,6 +669,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                output,
                encoded,
                kv_cache_dummy_dep=kv_cache_dummy_dep,
+                q_dcp_replicated=q_dcp_replicated,
            )
            return output

@@ -682,6 +687,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        quant_scale_ue8m0: bool | None = None,
        quant_col_major: bool | None = None,
        quant_tma_aligned: bool | None = None,
+        q_dcp_replicated: torch.Tensor | None = None,
    ) -> torch.Tensor:
        assert output is not None, "Output tensor must be provided."

@@ -738,6 +744,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
        output_padded = output
        output = output[:num_actual_toks, ...]
        q = q[:num_actual_toks, ...]
+        if q_dcp_replicated is not None:
+            q_dcp_replicated = q_dcp_replicated[:num_actual_toks, ...]
        k_c_normed = k_c_normed[:num_actual_toks, ...]
        k_pe = k_pe[:num_actual_toks, ...]

@@ -793,7 +801,12 @@ class MLAAttention(nn.Module, AttentionLayerBase):
            )

        if num_mqa_tokens > 0:
-            mqa_q = q[:num_mqa_tokens]
+            if q_dcp_replicated is not None:
+                mqa_q = q_dcp_replicated[:num_mqa_tokens]
+                qrep_decode = True
+            else:
+                mqa_q = q[:num_mqa_tokens]
+                qrep_decode = False
            mqa_output_slice = output[:num_mqa_tokens]

            mqa_q_nope, mqa_q_pe = mqa_q.split(
@@ -833,7 +846,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
            else:
                # Pads the head_dim if necessary (for the underlying kernel)
                N, B, P = mqa_q_nope.shape
-                _, _, L = self.W_UK_T.shape
+                W_UK_T = self.W_UK_T_dcp_qrep if qrep_decode else self.W_UK_T
+                assert W_UK_T is not None
+                _, _, L = W_UK_T.shape

                if self.q_pad_num_heads is not None:
                    mqa_ql_nope = mqa_q_nope.new_empty((self.q_pad_num_heads, B, L))
@@ -842,7 +857,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                    mqa_ql_nope = mqa_q_nope.new_empty((N, B, L))

                # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
-                torch.bmm(mqa_q_nope, self.W_UK_T, out=mqa_ql_nope)
+                torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope)

                # Convert from (N, B, L) to (B, N, L)
                mqa_ql_nope = mqa_ql_nope.transpose(0, 1)
@@ -855,6 +870,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                )
            else:
                mqa_q = (mqa_ql_nope, mqa_q_pe)
+            # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused)
            if self.impl.dcp_world_size > 1:
                if self.use_pcp:
                    if self.impl.dcp_world_size > self.impl.pcp_world_size:
@@ -863,8 +879,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                        mqa_q = get_tp_group().all_gather(mqa_q, dim=1)
                else:
                    if isinstance(mqa_q, tuple):
+                        # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
                        mqa_q = torch.cat(mqa_q, dim=-1)
-                    mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
+                    if not qrep_decode:
+                        # mqa_q do allgather in head dim.
+                        mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)

            # call decode attn
            if not self.impl.is_sparse:
@@ -952,6 +971,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):
            self.kv_b_proj, out_dtype=act_dtype
        ).T

+        if self.dcp_q_replicate:
+            # qrep wired here: validate unsupported decode backends once.
+            assert self.q_pad_num_heads in (None, self.num_heads), (
+                "DCP query replication is unsupported on head-padding MLA "
+                "backends (q_pad_num_heads)."
+            )
+            if (
+                self.is_aiter_triton_fp4_bmm_enabled
+                or self.is_aiter_triton_fp8_bmm_enabled
+            ):
+                raise NotImplementedError(
+                    "DCP query replication is not implemented for the aiter "
+                    "FP4/FP8 MLA BMM paths."
+                )
+
        assert kv_b_proj_weight.shape == (
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
@@ -1032,6 +1066,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
            replace_parameter(self, "W_UV", W_UV.transpose(0, 1), prefer_copy=True)
            # Convert from (L, N, P) to (N, P, L)
            replace_parameter(self, "W_UK_T", W_UK.permute(1, 2, 0), prefer_copy=True)
+            if self.dcp_q_replicate:
+                self.W_UK_T_dcp_qrep = get_dcp_group().all_gather(
+                    self.W_UK_T.contiguous(), dim=0
+                )

        # If we should not load quant weights, we initialize the scales to 1.0
        # as the default value. See [Note: Register q/k/v/prob scales in state dict]
@@ -1176,6 +1214,7 @@ def unified_mla_attention_with_output(
    quant_scale_ue8m0: bool | None = None,
    quant_col_major: bool | None = None,
    quant_tma_aligned: bool | None = None,
+    q_dcp_replicated: torch.Tensor | None = None,
 ) -> None:
    # kv_cache_dummy_dep is not used but accepting it creates a data dependency
    # that ensures torch.compile preserves ordering between KV cache update and
@@ -1196,6 +1235,7 @@ def unified_mla_attention_with_output(
        quant_scale_ue8m0=quant_scale_ue8m0,
        quant_col_major=quant_col_major,
        quant_tma_aligned=quant_tma_aligned,
+        q_dcp_replicated=q_dcp_replicated,
    )


@@ -1212,6 +1252,7 @@ def unified_mla_attention_with_output_fake(
    quant_scale_ue8m0: bool | None = None,
    quant_col_major: bool | None = None,
    quant_tma_aligned: bool | None = None,
+    q_dcp_replicated: torch.Tensor | None = None,
 ) -> None:
    return

@@ -34,6 +34,9 @@ class TrtLlmBf16ExpertsBase:
    monolithic interfaces.
    """

+    def supports_routing_replay_capture(self) -> bool:
+        return True
+
    def __init__(
        self,
        moe_config: FusedMoEConfig,
@@ -246,7 +249,11 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol

        assert activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]

-        return flashinfer.fused_moe.trtllm_bf16_moe(
+        routing_replay_out = self._maybe_make_routing_replay_buffer(
+            num_tokens=hidden_states.shape[0],
+            device=hidden_states.device,
+        )
+        out = flashinfer.fused_moe.trtllm_bf16_moe(
            routing_logits=router_logits,
            routing_bias=e_score_correction_bias,
            hidden_states=hidden_states,
@@ -263,4 +270,9 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol
            routing_method_type=self.routing_method_type,
            activation_type=activation_to_flashinfer_int(activation),
            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
+            routing_replay_out=routing_replay_out,
        )
+        self._maybe_dispatch_routing_replay(
+            routing_replay_out, num_tokens=hidden_states.shape[0]
+        )
+        return out
@@ -265,6 +265,9 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
    Fp8 TRTLLM-Gen MoE kernels. Supports monolithic interface.
    """

+    def supports_routing_replay_capture(self) -> bool:
+        return True
+
    def __init__(
        self,
        moe_config: FusedMoEConfig,
@@ -403,6 +406,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
            n_group = num_expert_group or 0
            selected_topk_group = topk_group or 0

+        routing_replay_out = self._maybe_make_routing_replay_buffer(
+            num_tokens=hidden_states.shape[0],
+            device=hidden_states.device,
+        )
+
        kwargs = dict(
            routing_logits=router_logits,
            routing_bias=e_score_correction_bias,
@@ -427,11 +435,16 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
            use_shuffled_weight=use_shuffled_weight,
            weight_layout=weight_layout,
            fp8_quantization_type=fp8_quant_type,
+            routing_replay_out=routing_replay_out,
            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
        )
        if is_mxfp8 or activation == MoEActivation.RELU2_NO_MUL:
            kwargs["activation_type"] = activation_type
-        return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs)
+        result = flashinfer.fused_moe.trtllm_fp8_block_scale_moe(**kwargs)
+        self._maybe_dispatch_routing_replay(
+            routing_replay_out, num_tokens=hidden_states.shape[0]
+        )
+        return result

    def _apply_per_tensor(
        self,
@@ -464,6 +477,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
        else:
            assert not apply_router_weight_on_input

+        routing_replay_out = self._maybe_make_routing_replay_buffer(
+            num_tokens=hidden_states.shape[0],
+            device=hidden_states.device,
+        )
+
        out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe(
            routing_logits=router_logits,
            routing_bias=e_score_correction_bias,
@@ -485,6 +503,10 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit
            routing_method_type=self.routing_method_type,
            activation_type=activation_type,
            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
+            routing_replay_out=routing_replay_out,
+        )
+        self._maybe_dispatch_routing_replay(
+            routing_replay_out, num_tokens=hidden_states.shape[0]
        )
        return out

@@ -122,6 +122,9 @@ class TrtLlmMxfp4ExpertsMonolithic(
    Wraps flashinfer.trtllm_fp4_block_scale_moe().
    """

+    def supports_routing_replay_capture(self) -> bool:
+        return True
+
    @staticmethod
    def _supports_parallel_config(
        moe_parallel_config: FusedMoEParallelConfig,
@@ -184,6 +187,10 @@ class TrtLlmMxfp4ExpertsMonolithic(
            device=hidden_states.device,
        )

+        routing_replay_out = self._maybe_make_routing_replay_buffer(
+            num_tokens=hidden_states.shape[0],
+            device=hidden_states.device,
+        )
        trtllm_fp4_block_scale_moe(
            routing_logits=router_logits.to(torch.bfloat16),
            routing_bias=None,
@@ -213,8 +220,11 @@ class TrtLlmMxfp4ExpertsMonolithic(
            do_finalize=True,
            tune_max_num_tokens=max(self.max_capture_size, 1),
            output=output,
+            routing_replay_out=routing_replay_out,
+        )
+        self._maybe_dispatch_routing_replay(
+            routing_replay_out, num_tokens=hidden_states.shape[0]
        )
-
        return output


@@ -122,6 +122,9 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
        # The kernel handles quantization internally.
        return True

+    def supports_routing_replay_capture(self) -> bool:
+        return True
+
    def apply(
        self,
        hidden_states: torch.Tensor,
@@ -144,7 +147,12 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):

        assert self.w1_scale is not None
        assert self.w2_scale is not None
-        return flashinfer_trtllm_mxint4_moe(
+
+        routing_replay_out = self._maybe_make_routing_replay_buffer(
+            num_tokens=hidden_states.shape[0],
+            device=hidden_states.device,
+        )
+        result = flashinfer_trtllm_mxint4_moe(
            x=hidden_states,
            router_logits=router_logits,
            w13_weight_packed=w1,
@@ -160,4 +168,9 @@ class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
            topk_group=topk_group,
            e_score_correction_bias=e_score_correction_bias,
            routing_method_type=self.routing_method,
+            routing_replay_out=routing_replay_out,
        )
+        self._maybe_dispatch_routing_replay(
+            routing_replay_out, num_tokens=hidden_states.shape[0]
+        )
+        return result
@@ -435,6 +435,9 @@ class TrtLlmNvFp4ExpertsMonolithic(
    Monolithic version of the kernel (router + experts).
    """

+    def supports_routing_replay_capture(self) -> bool:
+        return True
+
    @staticmethod
    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
        """The modular implementation should be used for the Dp/Ep or EPLB case."""
@@ -509,10 +512,14 @@ class TrtLlmNvFp4ExpertsMonolithic(

        output1_scale_gate_scalar = self.quant_config.g1_alphas

+        routing_replay_out = self._maybe_make_routing_replay_buffer(
+            num_tokens=hidden_states.shape[0],
+            device=hidden_states.device,
+        )
        # Invoke kernel.
        # NOTE: Activation padding and output
        # truncation are handled by the MoE runner's
-        return flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
+        result = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
            routing_logits=router_logits,
            routing_bias=e_score_correction_bias,
            hidden_states=hidden_states,
@@ -544,4 +551,9 @@ class TrtLlmNvFp4ExpertsMonolithic(
            activation_type=activation_to_flashinfer_int(activation),
            per_token_scale=per_token_scale,
            tune_max_num_tokens=fi_moe_largest_bucket(self.moe_config),
+            routing_replay_out=routing_replay_out,
        )[0]
+        self._maybe_dispatch_routing_replay(
+            routing_replay_out, num_tokens=hidden_states.shape[0]
+        )
+        return result
@@ -1000,6 +1000,57 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts):
    def is_monolithic() -> bool:
        return True

+    routing_replay_capture_fn: Callable[[torch.Tensor], None] | None = None
+    _routing_replay_buffer: torch.Tensor | None = None
+
+    def supports_routing_replay_capture(self) -> bool:
+        """Whether this expert supports routing replay capture.
+
+        Subclasses backed by a kernel that exposes routed expert IDs
+        (e.g. FlashInfer's ``routing_replay_out``) should override.
+        """
+        return False
+
+    def set_capture_fn(
+        self,
+        capture_fn: Callable[[torch.Tensor], None] | None,
+    ) -> None:
+        self.routing_replay_capture_fn = capture_fn
+        if capture_fn is None:
+            self._routing_replay_buffer = None
+            return
+        self._routing_replay_buffer = torch.empty(
+            (self.moe_config.max_num_tokens, self.moe_config.experts_per_token),
+            dtype=torch.int16,
+            device=self.moe_config.device,
+        )
+
+    def _maybe_make_routing_replay_buffer(
+        self,
+        num_tokens: int,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if self.routing_replay_capture_fn is None:
+            return None
+        buf = self._routing_replay_buffer
+        assert buf is not None
+        if buf.shape[0] < num_tokens or buf.device != device:
+            raise ValueError(
+                "Routing replay buffer was initialized for "
+                f"{buf.shape[0]} tokens on {buf.device}, but the kernel "
+                f"received {num_tokens} tokens on {device}."
+            )
+        return buf
+
+    def _maybe_dispatch_routing_replay(
+        self,
+        routing_replay_out: torch.Tensor | None,
+        num_tokens: int,
+    ) -> None:
+        if routing_replay_out is None or self.routing_replay_capture_fn is None:
+            return
+        self.routing_replay_capture_fn(routing_replay_out[:num_tokens])
+
    def apply(
        self,
        hidden_states: torch.Tensor,
@@ -90,10 +90,20 @@ class RoutedExpertsCapturer:
    ) -> None:
        hf_config = vllm_config.model_config.hf_text_config
        num_experts_per_tok = _get_num_experts_per_tok(hf_config)
+        num_layers = hf_config.num_hidden_layers
+        logger.info(
+            "RoutedExpertsCapturer: allocating buffer with "
+            "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d "
+            "(hf_config.model_type=%s)",
+            max_num_batched_tokens,
+            num_layers,
+            num_experts_per_tok,
+            getattr(hf_config, "model_type", "unknown"),
+        )
        self.device_buffer = torch.zeros(
            (
                max_num_batched_tokens,
-                hf_config.num_hidden_layers,
+                num_layers,
                num_experts_per_tok,
            ),
            # Use int32 for the device / host transit buffers: it
@@ -11,6 +11,7 @@ from torch.nn.parameter import Parameter
 from typing_extensions import TypeIs

 import vllm.envs as envs
+from vllm.config import get_current_vllm_config
 from vllm.distributed import (
    divide,
    get_tensor_model_parallel_rank,
@@ -254,6 +255,8 @@ class LinearBase(PluggableLayer):
        *,
        return_bias: bool = True,
        disable_tp: bool = False,
+        tp_rank: int | None = None,
+        tp_size: int | None = None,
    ):
        super().__init__()

@@ -277,8 +280,17 @@ class LinearBase(PluggableLayer):
            raise ValueError("All linear layers should support quant method.")
        self.return_bias = return_bias
        self.disable_tp = disable_tp
-        self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0
-        self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
+        if disable_tp:
+            self.tp_rank, self.tp_size = 0, 1
+        else:
+            self.tp_rank = (
+                tp_rank if tp_rank is not None else get_tensor_model_parallel_rank()
+            )
+            self.tp_size = (
+                tp_size
+                if tp_size is not None
+                else get_tensor_model_parallel_world_size()
+            )

    def update_param_tp_status(self):
        # Single source of truth for a parameter's TP state. BasevLLMParameter
@@ -425,6 +437,11 @@ class ColumnParallelLinear(LinearBase):
                        (e.g. model.layers.0.qkv_proj)
        return_bias: If true, return bias together with outputs in forward pass.
        disable_tp: If true, weights matrix won't be sharded through tp rank.
+        tp_rank: Override the tensor-parallel rank used for sharding. Defaults to
+            the global TP rank. Used to shard at a coarser granularity than one
+            shard per rank (see ``DCPGroupColumnParallelLinear``).
+        tp_size: Override the tensor-parallel world size used for sharding.
+            Defaults to the global TP world size.
    """

    # --8<-- [end:column_parallel_linear]
@@ -442,10 +459,21 @@ class ColumnParallelLinear(LinearBase):
        *,
        return_bias: bool = True,
        disable_tp: bool = False,
+        tp_rank: int | None = None,
+        tp_size: int | None = None,
    ):
        # Divide the weight matrix along the last dimension.
-        self.tp_rank = get_tensor_model_parallel_rank() if not disable_tp else 0
-        self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
+        if disable_tp:
+            self.tp_rank, self.tp_size = 0, 1
+        else:
+            self.tp_rank = (
+                tp_rank if tp_rank is not None else get_tensor_model_parallel_rank()
+            )
+            self.tp_size = (
+                tp_size
+                if tp_size is not None
+                else get_tensor_model_parallel_world_size()
+            )
        self.input_size_per_partition = input_size
        self.output_size_per_partition = divide(output_size, self.tp_size)
        self.output_partition_sizes = [self.output_size_per_partition]
@@ -465,6 +493,8 @@ class ColumnParallelLinear(LinearBase):
            prefix,
            return_bias=return_bias,
            disable_tp=disable_tp,
+            tp_rank=self.tp_rank,
+            tp_size=self.tp_size,
        )

        self._maybe_allow_fp8_block_shape_mismatch()
@@ -586,6 +616,47 @@ class ColumnParallelLinear(LinearBase):
        return s


+class DCPGroupColumnParallelLinear(ColumnParallelLinear):
+    """Column-parallel linear whose weight is sharded across DCP groups.
+
+    With Decode Context Parallelism (DCP) the KV cache is sharded across a DCP
+    group, so MLA decode must attend the group's full head set. This layer shards
+    its output across DCP *groups* (effective tp size ``tp_size //
+    dcp_world_size``) rather than across every rank, so each rank in a group
+    holds the whole group's heads, letting decode skip the query all-gather.
+
+    :meth:`forward` returns the group's full head set. :meth:`_local_view`
+    extracts this rank's TP shard for prefill.
+    """
+
+    def __init__(self, *args, **kwargs):
+        dcp_world_size = (
+            get_current_vllm_config().parallel_config.decode_context_parallel_size
+        )
+        rank = get_tensor_model_parallel_rank()
+        world_size = get_tensor_model_parallel_world_size()
+        self.group_size = max(dcp_world_size, 1)
+        self.qrep_active = self.group_size > 1
+        self.rank_in_group = rank % self.group_size
+        super().__init__(
+            *args,
+            **kwargs,
+            tp_rank=rank // self.group_size,
+            tp_size=world_size // self.group_size,
+        )
+
+    def _local_view(self, out: torch.Tensor) -> torch.Tensor:
+        """Slice this rank's tp head shard from a group-heads output.
+
+        ``out`` is head-shaped, i.e. ``(..., group_heads, head_dim)``.
+        """
+        if self.group_size == 1:
+            return out
+        n = out.shape[-2] // self.group_size
+        start = self.rank_in_group * n
+        return out[..., start : start + n, :].contiguous()
+
+
 class MergedColumnParallelLinear(ColumnParallelLinear):
    """Packed linear layers with column parallelism.

@@ -175,7 +175,7 @@ class MHCPreOp(CustomOp):
        norm_weight: torch.Tensor | None = None,
        norm_eps: float = 0.0,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        return self.forward_native(
+        return torch.ops._xpu_C.mhc_pre(
            residual,
            fn,
            hc_scale,
@@ -185,9 +185,6 @@ class MHCPreOp(CustomOp):
            hc_sinkhorn_eps,
            hc_post_mult_value,
            sinkhorn_repeat,
-            n_splits,
-            norm_weight,
-            norm_eps,
        )


@@ -260,7 +257,7 @@ class MHCPostOp(CustomOp):
        post_layer_mix: torch.Tensor,
        comb_res_mix: torch.Tensor,
    ) -> torch.Tensor:
-        return self.forward_native(
+        return torch.ops._xpu_C.mhc_post(
            x,
            residual,
            post_layer_mix,
@@ -369,16 +366,8 @@ class HCHeadOp(CustomOp):
        out = torch.empty(
            num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device
        )
-        torch.ops.vllm.hc_head_triton(
-            hs_flat,
-            hc_fn,
-            hc_scale,
-            hc_base,
-            out,
-            hidden_size,
-            rms_norm_eps,
-            hc_eps,
-            hc_mult,
+        torch.ops._xpu_C.hc_head_fused(
+            hs_flat, hc_fn, hc_scale, hc_base, out, rms_norm_eps, hc_eps
        )
        return out.view(*outer_shape, hidden_size)

@@ -548,7 +537,7 @@ class MHCFusedPostPreOp(CustomOp):
        norm_weight: torch.Tensor | None = None,
        norm_eps: float = 0.0,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        return self.forward_native(
+        return torch.ops._xpu_C.mhc_fused_post_pre(
            x,
            residual,
            post_layer_mix,
@@ -561,8 +550,4 @@ class MHCFusedPostPreOp(CustomOp):
            hc_sinkhorn_eps,
            hc_post_mult_value,
            sinkhorn_repeat,
-            n_splits,
-            tile_n,
-            norm_weight,
-            norm_eps,
        )
@@ -93,6 +93,10 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
        # the topk_tokens buffer written by a previous layer in the same pass.
        # Refer: https://arxiv.org/abs/2603.12201 for more details.
        self.skip_topk = skip_topk
+        # qrep is active when the query projection is a DCP-group-sharded layer
+        # that materializes the full group head set locally.
+        q_proj_layer = self.q_b_proj if self.q_lora_rank is not None else self.q_proj
+        self.dcp_q_replicate = getattr(q_proj_layer, "qrep_active", False)
        if self.indexer is not None:
            assert hasattr(self.indexer, "topk_tokens")
            self.topk_tokens = self.indexer.topk_tokens
@@ -110,6 +114,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
            quant_config=quant_config,
            prefix=f"{prefix}.attn",
            kv_b_proj=self.kv_b_proj,
+            dcp_q_replicate=self.dcp_q_replicate,
            use_sparse=self.is_sparse,
            indexer=self.indexer,
            topk_indices_buffer=mla_modules.topk_indices_buffer,
@@ -143,7 +148,8 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
                dim=-1,
            )
            q_c = self.q_a_layernorm(q_c)
-            q = self.q_b_proj(q_c)[0]
+            q_proj_layer = self.q_b_proj
+            q_proj_input = q_c
        else:
            assert self.kv_a_proj_with_mqa is not None, (
                "kv_a_proj_with_mqa is required when q_lora_rank is None"
@@ -152,15 +158,20 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
                "q_proj is required when q_lora_rank is None"
            )
            kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
-            q = self.q_proj(hidden_states)[0]
+            q_proj_layer = self.q_proj
+            q_proj_input = hidden_states

        kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        kv_c_normed = self.kv_a_layernorm(kv_c)
-
-        q = q.view(-1, self.num_heads, self.qk_head_dim)
        # Add head dim of 1 to k_pe
        k_pe = k_pe.unsqueeze(1)

+        q = q_proj_layer(q_proj_input)[0]
+        heads = self.num_heads
+        if self.dcp_q_replicate:
+            heads *= q_proj_layer.group_size
+        q = q.view(-1, heads, self.qk_head_dim)
+
        if self.rotary_emb is not None:
            q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
                positions, q[..., self.qk_nope_head_dim :], k_pe
@@ -172,11 +183,16 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
        if llama_4_scaling is not None:
            q *= llama_4_scaling

+        q_dcp_replicated = None
+        if self.dcp_q_replicate:
+            q_dcp_replicated, q = q, q_proj_layer._local_view(q)
+
        attn_out = self.mla_attn(
            q,
            kv_c_normed,
            k_pe,
            output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),
+            q_dcp_replicated=q_dcp_replicated,
        )

        return self.o_proj(attn_out)[0]
@@ -190,6 +190,7 @@ def flashinfer_trtllm_mxint4_moe(
    topk_group: int | None = None,
    e_score_correction_bias: torch.Tensor | None = None,
    routing_method_type: int | None = None,
+    routing_replay_out: torch.Tensor | None = None,
 ) -> torch.Tensor:
    """
    Apply FlashInfer TensorRT-LLM MxInt4 MoE kernel.
@@ -262,6 +263,7 @@ def flashinfer_trtllm_mxint4_moe(
        do_finalize=True,
        output=None,
        tune_max_num_tokens=8192,
+        routing_replay_out=routing_replay_out,
    )
    if isinstance(out, (tuple, list)):
        out = out[0]
@@ -37,6 +37,11 @@ class Cosmos3ForConditionalGeneration(Qwen3VLForConditionalGeneration):
            ".to_out.": ".o_proj.",
            ".norm_q.": ".q_norm.",
            ".norm_k.": ".k_norm.",
+            # ModelOpt-native dialect (diffusers/transformers read these; vLLM reads
+            # weight_scale/input_scale instead), drop so AutoWeightsLoader passes
+            ".input_quantizer.": None,
+            ".weight_quantizer.": None,
+            ".output_quantizer.": None,
        },
        orig_to_new_prefix={
            "proj_in.": None,
@@ -33,6 +33,7 @@ from torch import nn
 from transformers import DeepseekV2Config, DeepseekV3Config

 import vllm._custom_ops as ops
+import vllm.envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
@@ -56,6 +57,7 @@ from vllm.model_executor.layers.fused_moe import (
 from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
 from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
+    DCPGroupColumnParallelLinear,
    MergedColumnParallelLinear,
    QKVParallelLinear,
    ReplicatedLinear,
@@ -1015,9 +1017,17 @@ class DeepseekV2MLAAttention(nn.Module):
                prefix=f"{prefix}.kv_a_proj_with_mqa",
            )

+        qrep_enabled = (
+            envs.VLLM_DCP_Q_REPLICATE
+            and vllm_config.parallel_config.decode_context_parallel_size > 1
+            and vllm_config.parallel_config.prefill_context_parallel_size <= 1
+        )
+        q_proj_cls = (
+            DCPGroupColumnParallelLinear if qrep_enabled else ColumnParallelLinear
+        )
        if self.q_lora_rank is not None:
            self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps)
-            self.q_b_proj = ColumnParallelLinear(
+            self.q_b_proj = q_proj_cls(
                self.q_lora_rank,
                self.num_heads * self.qk_head_dim,
                bias=False,
@@ -1025,7 +1035,7 @@ class DeepseekV2MLAAttention(nn.Module):
                prefix=f"{prefix}.q_b_proj",
            )
        else:
-            self.q_proj = ColumnParallelLinear(
+            self.q_proj = q_proj_cls(
                proj_input_size,
                self.num_heads * self.qk_head_dim,
                bias=False,
@@ -84,6 +84,12 @@ from .utils import (

 logger = init_logger(__name__)

+_GEMMA4_EXPERT_PARENT_MAPPER = WeightsMapper(
+    orig_to_new_regex={
+        re.compile(r"(?<!\.moe)\.experts$"): ".moe.experts",
+    }
+)
+

 def _remap_gemma4_expert_weight_name(name: str) -> str:
    return re.sub(r"(?<!\.moe)\.experts\.(\d+)\.", r".moe.experts.\1.", name)
@@ -1508,7 +1514,7 @@ class Gemma4Model(nn.Module, EagleModelMixin):
 class Gemma4ForCausalLM(
    nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts, SupportsEagle3
 ):
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
        orig_to_new_prefix={
            # Gemma4ForConditionalGeneration already loads the text stack
            # from `model.language_model.*`. We reuse that same checkpoint
@@ -40,7 +40,10 @@ from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import ReplicatedLinear
-from vllm.model_executor.models.gemma4 import Gemma4ForCausalLM
+from vllm.model_executor.models.gemma4 import (
+    _GEMMA4_EXPERT_PARENT_MAPPER,
+    Gemma4ForCausalLM,
+)
 from vllm.model_executor.models.module_mapping import MultiModelKeys
 from vllm.model_executor.models.transformers.utils import recursive_replace_linear
 from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -998,7 +1001,7 @@ class Gemma4ForConditionalGeneration(
    }

    # Maps checkpoint prefixes to vLLM module paths.
-    hf_to_vllm_mapper = WeightsMapper(
+    hf_to_vllm_mapper = _GEMMA4_EXPERT_PARENT_MAPPER | WeightsMapper(
        orig_to_new_prefix={
            # vision tower
            "model.vision_tower": "vision_tower",
@@ -1010,7 +1013,7 @@ class Gemma4ForConditionalGeneration(
            "model.language_model.": "language_model.model.",
            "lm_head.": "language_model.lm_head.",
            "model": "language_model.model",
-        }
+        },
    )

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -651,9 +651,9 @@ class MiniCPMV4_6ProcessingInfo(MiniCPMVProcessingInfo):
 class MiniCPMV4_6ViTWindowAttentionSelfAttn(nn.Module):
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_stacked={
-            "q_proj": ("qkv_proj", "q"),
-            "k_proj": ("qkv_proj", "k"),
-            "v_proj": ("qkv_proj", "v"),
+            ".q_proj": (".qkv_proj", "q"),
+            ".k_proj": (".qkv_proj", "k"),
+            ".v_proj": (".qkv_proj", "v"),
        }
    )

@@ -2,6 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Run registered CuTeDSL warmup compile units."""

+# TODO(roberto): Remove this compatibility registry after registered CuTeDSL
+# warmups are migrated to the shared JIT warmup infrastructure.
+# https://github.com/vllm-project/vllm/pull/47451
+
 from __future__ import annotations

 import time
@@ -1,204 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""FA4 MLA prefill CuTeDSL compile warmup config."""
-
-from __future__ import annotations
-
-from collections.abc import Hashable, Iterator
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal
-
-import torch
-
-if TYPE_CHECKING:
-    from vllm.v1.attention.backends.fa_utils import (
-        FlashAttentionCuTeDSLCompileSpec,
-    )
-
-FA4ArchitectureFamily = Literal["sm90", "sm100f", "sm120"]
-
-FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16)
-
-# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan
-# covers qhead_per_kvhead=1.
-# Batch is not a current FA4 MLA-prefill key field. Use b1 for compile-only
-# specs because it is the conservative case for Split-KV shape heuristics.
-# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but vLLM
-# does not use them in this backend yet; they need a separate
-# num_kv_heads/qv/top-k-aware warmup plan if wired in later.
-FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1
-FA4_MLA_PREFILL_Q_TILE = 128
-FA4_MLA_PREFILL_K_TILE = 128
-FA4_MLA_PREFILL_LONG_K_BLOCKS = 32
-FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64
-FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True)
-FA4_MLA_PREFILL_LSE_OPTIONS = (False, True)
-
-
-@dataclass(frozen=True)
-class FA4MLAPrefillCompileContext:
-    dtype: torch.dtype
-    num_heads: int
-    qk_head_dim: int
-    v_head_dim: int
-    kv_nope_head_dim: int
-    requires_v_padding: bool
-    scale: float
-    num_splits: int
-    fa_version: int
-
-    # Return the V head dim FA4 sees.
-    @property
-    def effective_v_head_dim(self) -> int:
-        if self.requires_v_padding:
-            return self.qk_head_dim
-        return self.v_head_dim
-
-
-@dataclass(frozen=True)
-class FA4MLAPrefillCompileRequest:
-    """One compile-only FA4 MLA prefill request."""
-
-    key: Hashable
-    compile_spec: FlashAttentionCuTeDSLCompileSpec
-
-    # Compile this request.
-    def compile(self) -> None:
-        self.compile_spec.compile()
-
-
-# Yield deduped compile requests.
-def iter_fa4_mla_prefill_compile_requests(
-    ctx: FA4MLAPrefillCompileContext,
-) -> Iterator[FA4MLAPrefillCompileRequest]:
-    """Yield compile requests for this fixed MLA backend.
-
-    FA4 dedupes duplicate atomic kernel selections in its own JIT cache.
-    """
-    seen: set[Hashable] = set()
-    for compile_spec in iter_fa4_mla_prefill_compile_specs(ctx):
-        key = compile_spec.request_key()
-        if key in seen:
-            continue
-        seen.add(key)
-        yield FA4MLAPrefillCompileRequest(
-            key=key,
-            compile_spec=compile_spec,
-        )
-
-
-# Build compile specs for this setup.
-def iter_fa4_mla_prefill_compile_specs(
-    ctx: FA4MLAPrefillCompileContext,
-) -> Iterator[FlashAttentionCuTeDSLCompileSpec]:
-    """Yield compile-only FA4 MLA prefill requests for this fixed setup."""
-
-    arch_family = _fa4_architecture_family_from_compute_capability(
-        *torch.cuda.get_device_capability()
-    )
-    if not _supports_fa4_mla_prefill(ctx, arch_family):
-        return
-
-    from vllm.v1.attention.backends.fa_utils import (
-        FlashAttentionCuTeDSLCompileSpec,
-    )
-
-    batch_size = FA4_MLA_PREFILL_COMPILE_BATCH_SIZE
-    v_stride = None
-    if not ctx.requires_v_padding:
-        v_stride = (
-            ctx.num_heads * ctx.kv_nope_head_dim,
-            ctx.kv_nope_head_dim,
-            1,
-        )
-
-    for _, max_seqlen_q, max_seqlen_k in _shape_probes_for_context(ctx, arch_family):
-        total_q_tokens = batch_size * max_seqlen_q
-        total_kv_tokens = batch_size * max_seqlen_k
-        for causal in FA4_MLA_PREFILL_CAUSAL_OPTIONS:
-            for return_lse in FA4_MLA_PREFILL_LSE_OPTIONS:
-                yield FlashAttentionCuTeDSLCompileSpec(
-                    q_shape=(total_q_tokens, ctx.num_heads, ctx.qk_head_dim),
-                    k_shape=(total_kv_tokens, ctx.num_heads, ctx.qk_head_dim),
-                    v_shape=(
-                        total_kv_tokens,
-                        ctx.num_heads,
-                        ctx.effective_v_head_dim,
-                    ),
-                    v_stride=v_stride,
-                    q_dtype=ctx.dtype,
-                    cu_seqlens_q_shape=(batch_size + 1,),
-                    cu_seqlens_k_shape=(batch_size + 1,),
-                    max_seqlen_q=max_seqlen_q,
-                    max_seqlen_k=max_seqlen_k,
-                    softmax_scale=ctx.scale,
-                    causal=causal,
-                    return_softmax_lse=return_lse,
-                    num_splits=ctx.num_splits,
-                    fa_version=ctx.fa_version,
-                )
-
-
-# Pick one q/k point per current FA4 MLA-prefill shape regime.
-def _shape_probes_for_context(
-    ctx: FA4MLAPrefillCompileContext,
-    arch_family: FA4ArchitectureFamily,
-) -> tuple[tuple[str, int, int], ...]:
-    q_stage1_q = 1
-    q_stage2_q = FA4_MLA_PREFILL_Q_TILE + 1
-    # FA4 never auto-splits when ceil(max_seqlen_k / tile_n) <= 4.
-    no_split_k = 4 * FA4_MLA_PREFILL_K_TILE
-    long_k = FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
-    # Diff-head-dim Blackwell Split-KV switches tile_n at 64 K blocks.
-    very_long_k = FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE
-
-    base_probes = (
-        ("q_stage1", q_stage1_q, FA4_MLA_PREFILL_K_TILE),
-        ("q_stage2", q_stage2_q, no_split_k),
-    )
-    # SM120 currently rejects Split-KV in FA4; num_splits=1 also has no split
-    # shape regimes on any architecture.
-    if ctx.num_splits == 1 or arch_family == "sm120":
-        return base_probes
-
-    long_k_probes = (
-        ("q_stage1_long_k", q_stage1_q, long_k),
-        ("q_stage2_long_k", q_stage2_q, long_k),
-    )
-
-    # SM90 does not have the SM100 q_stage or diff-head-dim tile_n=64 branch.
-    # Same-dim SM100-family MLA also does not need the very-long-K probe.
-    if arch_family == "sm90" or ctx.qk_head_dim == ctx.effective_v_head_dim:
-        return (*base_probes, *long_k_probes)
-
-    very_long_k_probes = (
-        ("q_stage1_very_long_k", q_stage1_q, very_long_k),
-        ("q_stage2_very_long_k", q_stage2_q, very_long_k),
-    )
-    return (*base_probes, *long_k_probes, *very_long_k_probes)
-
-
-# Check whether this setup can use FA4 MLA prefill.
-def _supports_fa4_mla_prefill(
-    ctx: FA4MLAPrefillCompileContext,
-    arch_family: FA4ArchitectureFamily,
-) -> bool:
-    return (
-        ctx.dtype in FA4_STANDARD_DTYPES
-        and ctx.num_heads > 0
-        and (arch_family != "sm120" or ctx.num_splits == 1)
-    )
-
-
-# Map CUDA capability to the FA4 arch family used by warmup checks.
-def _fa4_architecture_family_from_compute_capability(
-    major: int,
-    minor: int,
-) -> FA4ArchitectureFamily:
-    if (major, minor) == (9, 0):
-        return "sm90"
-    if major == 10:
-        return "sm100f"
-    if (major, minor) == (12, 0):
-        return "sm120"
-    raise ValueError(f"FA4 warmup does not know CUDA capability {major}.{minor}")
@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Warm up FA4 CuTeDSL MLA prefill compile keys."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu_worker import Worker
+
+
+def fa4_cutedsl_warmup(worker: Worker) -> None:
+    runner = worker.model_runner
+    if runner.is_pooling_model:
+        return
+
+    vllm_config = runner.vllm_config
+    if not vllm_config.model_config.use_mla:
+        return
+
+    backend_cls = get_mla_prefill_backend(vllm_config)
+    if backend_cls.get_name() != "FLASH_ATTN":
+        return
+
+    from vllm.v1.attention.backends.mla.prefill import flash_attn
+
+    flash_attn.FA4_MLA_PREFILL_KERNEL.warmup(vllm_config)
@@ -0,0 +1,474 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared interfaces and tracing helpers for explicit JIT warmup keys."""
+
+from __future__ import annotations
+
+import ast
+import inspect
+import itertools
+import operator
+import textwrap
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable, Mapping
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+__all__ = [
+    "VllmJitKernel",
+    "WarmupIntRange",
+    "get_ast_full_name",
+    "get_function_source_node",
+    "zip_inputs",
+]
+
+
+CompileKeyT = TypeVar("CompileKeyT")
+
+
+@dataclass(frozen=True)
+class WarmupIntRange:
+    start: int
+    stop: int
+    step: int = 1
+
+
+WarmupValues = Any
+CompileKeyDispatchFn = Callable[..., CompileKeyT]
+
+
+@dataclass(frozen=True)
+class _WarmupInputRows:
+    """Warmup dispatch inputs expanded in lockstep."""
+
+    rows: tuple[Mapping[str, WarmupValues], ...]
+
+
+def _expand_warmup_values(values: WarmupValues) -> tuple[Any, ...]:
+    if isinstance(values, WarmupIntRange):
+        return tuple(range(values.start, values.stop, values.step))
+    if isinstance(values, (list, tuple)):
+        return tuple(values)
+    return (values,)
+
+
+def zip_inputs(*rows: Mapping[str, WarmupValues]) -> _WarmupInputRows:
+    """Group row-wise dispatch inputs that should be expanded in lockstep."""
+    if not rows:
+        raise ValueError("zip_inputs requires at least one dispatch input row")
+    if not all(isinstance(row, Mapping) for row in rows):
+        raise ValueError("zip_inputs rows must be mappings")
+
+    first_names = frozenset(rows[0])
+    if not first_names:
+        raise ValueError("zip_inputs rows require at least one dispatch input name")
+    if not all(isinstance(name, str) for name in first_names):
+        raise ValueError("zip_inputs dispatch input names must be strings")
+
+    input_rows: list[Mapping[str, WarmupValues]] = []
+    for row in rows:
+        names = frozenset(row)
+        if names != first_names:
+            raise ValueError("zip_inputs rows must use the same dispatch input names")
+        input_rows.append(dict(row))
+
+    return _WarmupInputRows(rows=tuple(input_rows))
+
+
+def _expand_warmup_value_grid(
+    values: Mapping[str, WarmupValues],
+    input_names: frozenset[str],
+) -> tuple[dict[str, Any], ...]:
+    names = tuple(name for name in values if name in input_names)
+    if not names:
+        return ({},)
+
+    expanded_values = tuple(_expand_warmup_values(values[name]) for name in names)
+    return tuple(
+        dict(zip(names, value_set)) for value_set in itertools.product(*expanded_values)
+    )
+
+
+def _expand_warmup_input_rows(
+    rows: tuple[Mapping[str, WarmupValues], ...],
+    input_names: frozenset[str],
+) -> tuple[dict[str, Any], ...]:
+    active_names = frozenset(name for name in rows[0] if name in input_names)
+    if not active_names:
+        return ({},)
+
+    return tuple(
+        {name: value for name, value in row.items() if name in active_names}
+        for row in rows
+    )
+
+
+def _merge_warmup_kwargs(parts: Iterable[Mapping[str, Any]]) -> dict[str, Any]:
+    merged: dict[str, Any] = {}
+    for part in parts:
+        for name, value in part.items():
+            if name in merged:
+                raise ValueError(
+                    f"Warmup dispatch input '{name}' is specified more than once"
+                )
+            merged[name] = value
+    return merged
+
+
+@dataclass(frozen=True)
+class _CompileKeyDispatchTrace:
+    local_exprs: tuple[tuple[str, ast.AST], ...]
+    field_exprs: tuple[tuple[str, ast.AST], ...]
+    globals: Mapping[str, Any]
+    input_names: frozenset[str]
+    defaults: Mapping[str, Any]
+
+    def compile_key(
+        self,
+        compile_key_type: type[CompileKeyT],
+        kwargs: Mapping[str, Any],
+    ) -> CompileKeyT:
+        dispatch_values = {**self.defaults, **kwargs}
+        for name, expr in self.local_exprs:
+            dispatch_values[name] = _eval_dispatch_expr(
+                expr, dispatch_values, self.globals
+            )
+        return compile_key_type(
+            **{
+                field: _eval_dispatch_expr(expr, dispatch_values, self.globals)
+                for field, expr in self.field_exprs
+            }
+        )
+
+
+_BIN_OPS: dict[type[ast.operator], Callable[[Any, Any], Any]] = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.FloorDiv: operator.floordiv,
+    ast.Mod: operator.mod,
+    ast.Pow: operator.pow,
+}
+_CMP_OPS: dict[type[ast.cmpop], Callable[[Any, Any], bool]] = {
+    ast.Eq: operator.eq,
+    ast.NotEq: operator.ne,
+    ast.Lt: operator.lt,
+    ast.LtE: operator.le,
+    ast.Gt: operator.gt,
+    ast.GtE: operator.ge,
+}
+
+
+def _dispatch_expr_source(node: ast.AST) -> str:
+    try:
+        return ast.unparse(node)
+    except Exception:
+        return ast.dump(node)
+
+
+def _dispatch_expr_error(node: ast.AST, reason: str) -> ValueError:
+    return ValueError(
+        f"{reason}: {_dispatch_expr_source(node)}. "
+        "Supported dispatch expressions are names, constants, attributes, "
+        "tuple/list literals, conditional expressions, comparisons, boolean "
+        "operators, unary not/minus, arithmetic, and calls without **kwargs."
+    )
+
+
+class _DispatchExprEvaluator(ast.NodeVisitor):
+    def __init__(
+        self,
+        values: Mapping[str, Any],
+        globals_: Mapping[str, Any],
+    ) -> None:
+        self.values = values
+        self.globals = globals_
+
+    def eval(self, node: ast.AST) -> Any:
+        return self.visit(node)
+
+    def generic_visit(self, node: ast.AST) -> Any:
+        raise _dispatch_expr_error(node, "Unsupported dispatch expression")
+
+    def visit_Name(self, node: ast.Name) -> Any:
+        if node.id in self.values:
+            return self.values[node.id]
+        if node.id in self.globals:
+            return self.globals[node.id]
+        raise _dispatch_expr_error(node, f"Unknown dispatch name '{node.id}'")
+
+    def visit_Constant(self, node: ast.Constant) -> Any:
+        return node.value
+
+    def visit_IfExp(self, node: ast.IfExp) -> Any:
+        return self.visit(node.body if self.visit(node.test) else node.orelse)
+
+    def visit_Tuple(self, node: ast.Tuple) -> tuple[Any, ...]:
+        return tuple(self.visit(elt) for elt in node.elts)
+
+    def visit_List(self, node: ast.List) -> list[Any]:
+        return [self.visit(elt) for elt in node.elts]
+
+    def visit_BoolOp(self, node: ast.BoolOp) -> Any:
+        if isinstance(node.op, ast.And):
+            result = None
+            for value in node.values:
+                result = self.visit(value)
+                if not result:
+                    return result
+            return result
+        if isinstance(node.op, ast.Or):
+            result = None
+            for value in node.values:
+                result = self.visit(value)
+                if result:
+                    return result
+            return result
+        raise _dispatch_expr_error(node, "Unsupported dispatch boolean operator")
+
+    def visit_Compare(self, node: ast.Compare) -> bool:
+        left = self.visit(node.left)
+        for op_node, comparator in zip(node.ops, node.comparators):
+            right = self.visit(comparator)
+            op = _CMP_OPS.get(type(op_node))
+            if op is None:
+                raise _dispatch_expr_error(
+                    node, "Unsupported dispatch comparison operator"
+                )
+            if not op(left, right):
+                return False
+            left = right
+        return True
+
+    def visit_UnaryOp(self, node: ast.UnaryOp) -> Any:
+        operand = self.visit(node.operand)
+        if isinstance(node.op, ast.Not):
+            return not operand
+        if isinstance(node.op, ast.USub):
+            return -operand
+        raise _dispatch_expr_error(node, "Unsupported dispatch unary operator")
+
+    def visit_BinOp(self, node: ast.BinOp) -> Any:
+        op = _BIN_OPS.get(type(node.op))
+        if op is None:
+            raise _dispatch_expr_error(node, "Unsupported dispatch binary operator")
+        return op(self.visit(node.left), self.visit(node.right))
+
+    def visit_Call(self, node: ast.Call) -> Any:
+        args = [self.visit(arg) for arg in node.args]
+        fn = self.visit(node.func)
+        call_kwargs: dict[str, Any] = {}
+        for keyword in node.keywords:
+            if keyword.arg is None:
+                raise _dispatch_expr_error(
+                    node, "Dispatch helper calls cannot use **kwargs"
+                )
+            call_kwargs[keyword.arg] = self.visit(keyword.value)
+        return fn(*args, **call_kwargs)
+
+    def visit_Attribute(self, node: ast.Attribute) -> Any:
+        return getattr(self.visit(node.value), node.attr)
+
+
+def get_ast_full_name(node: ast.AST) -> str | None:
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        parent = get_ast_full_name(node.value)
+        if parent is not None:
+            return f"{parent}.{node.attr}"
+    return None
+
+
+def get_function_source_node(fn: Callable[..., Any]) -> ast.FunctionDef:
+    source_fn = getattr(fn, "fn", fn)
+    source = textwrap.dedent(inspect.getsource(source_fn))
+    tree = ast.parse(source)
+    function_defs = [node for node in tree.body if isinstance(node, ast.FunctionDef)]
+    if len(function_defs) != 1:
+        name = getattr(source_fn, "__name__", type(source_fn).__name__)
+        raise ValueError(f"Expected one function in {name}, found {len(function_defs)}")
+    return function_defs[0]
+
+
+def _eval_dispatch_expr(
+    node: ast.AST,
+    kwargs: Mapping[str, Any],
+    globals_: Mapping[str, Any],
+) -> Any:
+    return _DispatchExprEvaluator(kwargs, globals_).eval(node)
+
+
+def _collect_input_names(
+    node: ast.AST,
+    candidate_names: set[str],
+    local_names: set[str] | None = None,
+) -> set[str]:
+    if local_names is None:
+        local_names = set()
+    return {
+        child.id
+        for child in ast.walk(node)
+        if (
+            isinstance(child, ast.Name)
+            and child.id in candidate_names
+            and child.id not in local_names
+        )
+    }
+
+
+def _collect_dispatch_body(
+    fn: CompileKeyDispatchFn[Any],
+    function_def: ast.FunctionDef,
+) -> tuple[list[tuple[str, ast.AST]], ast.Call]:
+    local_exprs: list[tuple[str, ast.AST]] = []
+    for statement in function_def.body:
+        if (
+            isinstance(statement, ast.Expr)
+            and isinstance(statement.value, ast.Constant)
+            and isinstance(statement.value.value, str)
+        ):
+            continue
+
+        if isinstance(statement, ast.Assign):
+            if len(statement.targets) != 1 or not isinstance(
+                statement.targets[0], ast.Name
+            ):
+                raise _dispatch_expr_error(
+                    statement, "Dispatch assignments must target one local name"
+                )
+            local_exprs.append((statement.targets[0].id, statement.value))
+            continue
+
+        if isinstance(statement, ast.AnnAssign):
+            if statement.value is None:
+                raise _dispatch_expr_error(
+                    statement, "Dispatch annotations must assign a value"
+                )
+            if not isinstance(statement.target, ast.Name):
+                raise _dispatch_expr_error(
+                    statement, "Dispatch assignments must target one local name"
+                )
+            local_exprs.append((statement.target.id, statement.value))
+            continue
+
+        if isinstance(statement, ast.Return) and isinstance(statement.value, ast.Call):
+            return local_exprs, statement.value
+
+        if isinstance(statement, ast.Return):
+            raise _dispatch_expr_error(
+                statement, "Dispatch must return one CompileKey(...) call"
+            )
+
+        raise _dispatch_expr_error(
+            statement,
+            "Dispatch may only contain local assignments before CompileKey return",
+        )
+
+    raise ValueError(f"Expected {fn.__name__} to return one CompileKey(...) call")
+
+
+def _trace_compile_key_dispatch(
+    fn: CompileKeyDispatchFn[Any],
+) -> _CompileKeyDispatchTrace:
+    source_fn = getattr(fn, "__func__", fn)
+    globals_ = source_fn.__globals__
+    function_def = get_function_source_node(fn)
+
+    local_exprs, return_call = _collect_dispatch_body(fn, function_def)
+
+    field_exprs: list[tuple[str, ast.AST]] = []
+    signature = inspect.signature(fn)
+    defaults = {
+        name: parameter.default
+        for name, parameter in signature.parameters.items()
+        if parameter.default is not inspect.Parameter.empty
+    }
+    candidate_names = set(signature.parameters)
+    input_names: set[str] = set()
+    local_names = {name for name, _ in local_exprs}
+    for _, expr in local_exprs:
+        input_names.update(_collect_input_names(expr, candidate_names))
+    for keyword in return_call.keywords:
+        if keyword.arg is None:
+            raise ValueError(f"{fn.__name__} cannot use **kwargs in CompileKey")
+        field_exprs.append((keyword.arg, keyword.value))
+        input_names.update(
+            _collect_input_names(keyword.value, candidate_names, local_names)
+        )
+
+    return _CompileKeyDispatchTrace(
+        tuple(local_exprs),
+        tuple(field_exprs),
+        globals_,
+        frozenset(input_names),
+        defaults,
+    )
+
+
+class VllmJitKernel(Generic[CompileKeyT], ABC):
+    """Kernel wrapper that owns dispatch, warmup keys, and compilation."""
+
+    CompileKey: type[CompileKeyT]
+
+    def __init__(self) -> None:
+        self.compile_key_dispatch_trace = _trace_compile_key_dispatch(self.dispatch)
+
+    def compile_key(self, kwargs: Mapping[str, Any]) -> CompileKeyT:
+        return self.compile_key_dispatch_trace.compile_key(self.CompileKey, kwargs)
+
+    def _trace_dispatch(
+        self, dispatch: CompileKeyDispatchFn[CompileKeyT]
+    ) -> Callable[..., list[CompileKeyT]]:
+        compile_key_dispatch_trace = _trace_compile_key_dispatch(dispatch)
+
+        def traced(
+            *input_groups: _WarmupInputRows,
+            **kwargs: WarmupValues,
+        ) -> list[CompileKeyT]:
+            for group in input_groups:
+                if not isinstance(group, _WarmupInputRows):
+                    raise TypeError(
+                        "_trace_dispatch positional arguments must be "
+                        "zip_inputs(...) groups"
+                    )
+            expanded_input_groups = tuple(
+                _expand_warmup_input_rows(
+                    group.rows, compile_key_dispatch_trace.input_names
+                )
+                for group in input_groups
+            )
+            expanded_kwargs = _expand_warmup_value_grid(
+                kwargs, compile_key_dispatch_trace.input_names
+            )
+            dispatch_value_groups = (*expanded_input_groups, expanded_kwargs)
+            return list(
+                dict.fromkeys(
+                    compile_key_dispatch_trace.compile_key(
+                        self.CompileKey, _merge_warmup_kwargs(dispatch_values)
+                    )
+                    for dispatch_values in itertools.product(*dispatch_value_groups)
+                )
+            )
+
+        return traced
+
+    @abstractmethod
+    def dispatch(self, **kwargs: Any) -> CompileKeyT:
+        """Build one compile key from one concrete dispatch point."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_warmup_keys(self, *args: Any, **kwargs: Any) -> list[CompileKeyT]:
+        """Return compile keys that should be warmed for this kernel."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def compile(self, compile_key: CompileKeyT) -> None:
+        """Compile one warmup key."""
+        raise NotImplementedError
+
+    def warmup(self, *args: Any, **kwargs: Any) -> None:
+        """Compile this kernel's warmup keys."""
+        for compile_key in self.get_warmup_keys(*args, **kwargs):
+            self.compile(compile_key)
@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import ast
+import inspect
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from vllm.model_executor.warmup.jit_warmup import (
+    get_ast_full_name,
+    get_function_source_node,
+)
+
+
+@dataclass(frozen=True)
+class TritonWarmupTensor:
+    # Compile-only tensor descriptor for Triton pointer specialization.
+    dtype: Any
+    aligned: bool = True
+    shape: tuple[int, ...] = (1,)
+
+    def data_ptr(self) -> int:
+        return 0 if self.aligned else 1
+
+    def ptr_range(self) -> int:
+        return 0
+
+    def stride(self) -> tuple[int, ...]:
+        strides: list[int] = []
+        stride = 1
+        for size in reversed(self.shape):
+            strides.append(stride)
+            stride *= size
+        return tuple(reversed(strides))
+
+
+@dataclass(frozen=True)
+class TritonPointerInputVariant:
+    # Named pointer-alignment variant for compile-only Triton warmup.
+    alignments: tuple[tuple[str, bool], ...]
+
+    @classmethod
+    def from_alignment(cls, **aligned: bool) -> "TritonPointerInputVariant":
+        return cls(tuple(aligned.items()))
+
+    def is_aligned(self, name: str) -> bool:
+        for alignment_name, aligned in self.alignments:
+            if alignment_name == name:
+                return aligned
+        raise KeyError(f"Unknown Triton pointer input variant: {name}")
+
+    def pointer(
+        self,
+        name: str,
+        dtype: Any,
+        shape: tuple[int, ...] = (1,),
+    ) -> TritonWarmupTensor:
+        return TritonWarmupTensor(dtype, aligned=self.is_aligned(name), shape=shape)
+
+
+def _literal_str_refs(node: ast.AST) -> tuple[str | int, ...]:
+    if isinstance(node, ast.Constant) and isinstance(node.value, str | int):
+        return (node.value,)
+    if isinstance(node, ast.List | ast.Tuple):
+        refs: list[str | int] = []
+        for elt in node.elts:
+            if isinstance(elt, ast.Constant) and isinstance(elt.value, str | int):
+                refs.append(elt.value)
+            else:
+                raise ValueError(
+                    f"Unsupported Triton specialization ref: {ast.dump(elt)}"
+                )
+        return tuple(refs)
+    raise ValueError(f"Unsupported Triton specialization refs: {ast.dump(node)}")
+
+
+def _normalize_arg_refs(
+    refs: tuple[str | int, ...],
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    names: set[str] = set()
+    for ref in refs:
+        if isinstance(ref, int):
+            names.add(arg_names[ref])
+        else:
+            names.add(ref)
+    return frozenset(names)
+
+
+def _decorator_keyword_refs(
+    function_def: ast.FunctionDef,
+    keyword_name: str,
+) -> tuple[str | int, ...]:
+    for decorator in function_def.decorator_list:
+        if not isinstance(decorator, ast.Call):
+            continue
+        decorator_name = get_ast_full_name(decorator.func)
+        if decorator_name not in ("triton.jit", "jit"):
+            continue
+        for keyword in decorator.keywords:
+            if keyword.arg == keyword_name:
+                return _literal_str_refs(keyword.value)
+    return ()
+
+
+def _triton_do_not_specialize_args(
+    kernel: Callable[..., Any],
+    function_def: ast.FunctionDef,
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    refs = getattr(kernel, "do_not_specialize", None)
+    if refs is not None:
+        return _normalize_arg_refs(tuple(refs), arg_names)
+    return _normalize_arg_refs(
+        _decorator_keyword_refs(function_def, "do_not_specialize"),
+        arg_names,
+    )
+
+
+def _triton_constexpr_arg_names(
+    kernel: Callable[..., Any],
+    function_def: ast.FunctionDef,
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    constexprs = getattr(kernel, "constexprs", None)
+    if constexprs is not None:
+        return frozenset(arg_names[index] for index in constexprs)
+
+    names: set[str] = set()
+    for arg in function_def.args.args + function_def.args.kwonlyargs:
+        if arg.annotation is None:
+            continue
+        annotation = get_ast_full_name(arg.annotation)
+        if annotation in ("tl.constexpr", "triton.language.constexpr", "constexpr"):
+            names.add(arg.arg)
+    return frozenset(names)
+
+
+def _leftmost_name(node: ast.AST) -> str | None:
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.BinOp):
+        return _leftmost_name(node.left)
+    return None
+
+
+def _pointer_arg_names(
+    function_def: ast.FunctionDef,
+    arg_names: tuple[str, ...],
+) -> frozenset[str]:
+    candidate_names = set(arg_names)
+    pointer_names = {name for name in arg_names if name.endswith("_ptr")}
+    for node in ast.walk(function_def):
+        if not isinstance(node, ast.Call):
+            continue
+        if get_ast_full_name(node.func) not in ("tl.load", "tl.store"):
+            continue
+        if not node.args:
+            continue
+        name = _leftmost_name(node.args[0])
+        if name in candidate_names:
+            pointer_names.add(name)
+    return frozenset(pointer_names)
+
+
+def trace_triton_kernel_specialization_args(
+    kernel: Callable[..., Any],
+) -> tuple[str, ...]:
+    function_def = get_function_source_node(kernel)
+    source_fn = getattr(kernel, "fn", kernel)
+    arg_names = tuple(inspect.signature(source_fn).parameters)
+    constexpr_args = _triton_constexpr_arg_names(kernel, function_def, arg_names)
+    do_not_specialize_args = _triton_do_not_specialize_args(
+        kernel, function_def, arg_names
+    )
+    pointer_args = _pointer_arg_names(function_def, arg_names)
+
+    return tuple(
+        name
+        for name in arg_names
+        if name in constexpr_args
+        or (name not in pointer_args and name not in do_not_specialize_args)
+    )
@@ -17,6 +17,9 @@ from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
 from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
    deepseek_v4_mhc_warmup,
 )
+from vllm.model_executor.warmup.fa4_cutedsl_warmup import (
+    fa4_cutedsl_warmup,
+)
 from vllm.model_executor.warmup.flashinfer_autotune_cache import (
    resolve_flashinfer_autotune_file,
    write_flashinfer_autotune_cache,
@@ -27,7 +30,7 @@ from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import (
 )
 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
 from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
-    sparse_mla_triton_warmup_if_needed,
+    sparse_mla_triton_warmup,
 )
 from vllm.model_executor.warmup.v1_block_table_warmup import (
    warm_v1_block_table_kernels,
@@ -94,7 +97,6 @@ def kernel_warmup(worker: "Worker"):
    )

    # Run next so input-prep kernels JIT against pristine runner state.
-    sparse_mla_triton_warmup_if_needed(worker)
    flashinfer_sparse_mla_decode_autotune_warmup(worker)
    deepseek_v4_sparse_mla_attention_warmup(worker)

@@ -156,8 +158,15 @@ def kernel_warmup(worker: "Worker"):
        )

    if worker.vllm_config.kernel_config.enable_cutedsl_warmup:
+        # TODO(roberto): Remove after registered CuTeDSL warmups are migrated
+        # to the shared JIT warmup infrastructure.
+        # https://github.com/vllm-project/vllm/pull/47451
        cutedsl_warmup()

+    if worker.vllm_config.kernel_config.enable_jit_warmup:
+        fa4_cutedsl_warmup(worker)
+        sparse_mla_triton_warmup(worker)
+

 def _flashinfer_autotune_skip_ops(runner: "GPUModelRunner") -> set[str] | None:
    if envs.VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS is not None:
@@ -4,11 +4,10 @@

 from typing import TYPE_CHECKING

-import torch
-
 from vllm.logger import init_logger

 if TYPE_CHECKING:
+    from vllm.config import VllmConfig
    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
    from vllm.v1.worker.gpu_worker import Worker

@@ -31,47 +30,7 @@ _GENERIC_SPARSE_MLA_BACKENDS = frozenset(
 )
 _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})

-_SPARSE_PREFILL_METADATA_NUM_PREFILLS = (1, 2, 4, 8)
-_SPARSE_PREFILL_METADATA_NUM_DECODES = (0, 1, 2)
-_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS = (4, 128)
-_PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS = (2, 3)
-_PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS = (
-    # query_slice_start offset, query_slice_stop offset
-    (0, 0),
-    (0, -1),
-    (1, 0),
-    (1, -1),
-)
-_COMBINE_TOPK_SWA_INPUT_VARIANTS = (
-    # offset_topk, offset_query_and_seq, offset_gather
-    (False, False, False),
-    (False, True, False),
-    (True, True, True),
-)
-_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES = (
-    # compress_ratio, topk, topk_width, N
-    (1, 0, 512, 512),
-    (4, 512, 512, 512 * 4),
-    # DSv4-Pro C4A traffic uses top-k 1024 with N=1024.
-    (4, 1024, 1024, 1024),
-    (128, 8192, 8192, 8192 * 128),
-    # Real C128A traffic also specializes N=1 in one call path.
-    (128, 8192, 8192, 1),
-)
-
-
-def _clamp_warmup_tokens(num_tokens: int, max_tokens: int) -> int:
-    return max(0, min(num_tokens, max_tokens))
-
-
-def _next_power_of_2(x: int) -> int:
-    return 1 << (x - 1).bit_length()
-
-
-def _hf_config_int(runner: "GPUModelRunner", name: str, default: int) -> int:
-    model_config = getattr(runner.vllm_config, "model_config", None)
-    hf_config = getattr(model_config, "hf_config", None)
-    return int(getattr(hf_config, name, default) or default)
+_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})


 def _attention_backend_name(backend: object) -> str | None:
@@ -96,242 +55,57 @@ def _has_attention_backend(
    return False


-def _warm_sparse_swa_prefill_metadata_kernel(
-    device: torch.device,
-    window_size: int,
-    prefill_tokens: int,
+def _compile_sparse_swa_prefill_metadata_kernel(
+    vllm_config: "VllmConfig",
 ) -> None:
    from vllm.v1.attention.backends.mla.sparse_swa import (
-        _compute_prefill_metadata_kernel,
+        _COMPUTE_PREFILL_METADATA_KERNEL,
    )

-    for num_prefills in _SPARSE_PREFILL_METADATA_NUM_PREFILLS:
-        for num_decodes in _SPARSE_PREFILL_METADATA_NUM_DECODES:
-            query_lens = [1] * num_decodes
-            query_lens += [prefill_tokens] * num_prefills
-            query_start_locs = [0]
-            for query_len in query_lens:
-                query_start_locs.append(query_start_locs[-1] + query_len)
-            query_start_loc = torch.tensor(
-                query_start_locs,
-                dtype=torch.int32,
-                device=device,
-            )
-            seq_lens = torch.tensor(
-                [1] * num_decodes + [window_size + q for q in query_lens[num_decodes:]],
-                dtype=torch.int32,
-                device=device,
-            )
-            prefill_gather_lens = torch.empty(
-                num_prefills, dtype=torch.int32, device=device
-            )
-            _compute_prefill_metadata_kernel[(1,)](
-                prefill_gather_lens,
-                seq_lens,
-                query_start_loc,
-                num_prefills,
-                num_decodes,
-                window_size,
-                BLOCK_SIZE=_next_power_of_2(num_prefills),
-            )
+    _COMPUTE_PREFILL_METADATA_KERNEL.warmup(vllm_config)


-def _warm_prefill_chunk_metadata_kernel(
-    device: torch.device,
-    compress_ratio: int,
-    query_len: int,
+def _compile_prefill_chunk_metadata_kernel(
+    vllm_config: "VllmConfig",
 ) -> None:
-    from vllm.v1.attention.backends.mla.indexer import build_prefill_chunk_metadata
-
-    num_reqs = 2
-    query_start_loc_cpu = torch.arange(
-        0, (num_reqs + 1) * query_len, query_len, dtype=torch.int32
-    )
-    query_start_loc = query_start_loc_cpu.to(device=device)
-
-    uncompressed_seq_lens_cpu = torch.tensor(
-        [
-            compress_ratio * multiplier + query_len
-            for multiplier in _PREFILL_CHUNK_METADATA_SEQ_LEN_MULTIPLIERS
-        ],
-        dtype=torch.int32,
-    )
-    compressed_seq_lens_cpu = uncompressed_seq_lens_cpu // compress_ratio
-    uncompressed_seq_lens = uncompressed_seq_lens_cpu.to(device=device)
-    compressed_seq_lens = compressed_seq_lens_cpu.to(device=device)
-    block_table = torch.zeros(
-        (num_reqs, int(compressed_seq_lens_cpu.max().item())),
-        dtype=torch.int32,
-        device=device,
+    from vllm.v1.attention.backends.mla.indexer import (
+        _BUILD_PREFILL_CHUNK_METADATA_KERNEL,
    )

-    offset_uncompressed_seq_lens = torch.empty(
-        num_reqs + 1, dtype=torch.int32, device=device
-    )[1:]
-    offset_uncompressed_seq_lens.copy_(uncompressed_seq_lens)
-    query_slices = tuple(
-        slice(start, num_reqs * query_len + stop)
-        for start, stop in _PREFILL_CHUNK_METADATA_QUERY_SLICE_OFFSETS
-    )
-    for warmup_uncompressed_seq_lens in (
-        uncompressed_seq_lens,
-        offset_uncompressed_seq_lens,
-    ):
-        for query_slice in query_slices:
-            build_prefill_chunk_metadata(
-                0,
-                num_reqs,
-                query_start_loc,
-                query_start_loc_cpu,
-                warmup_uncompressed_seq_lens,
-                compressed_seq_lens,
-                compressed_seq_lens_cpu,
-                block_table,
-                compress_ratio,
-                query_slice=query_slice,
-            )
+    _BUILD_PREFILL_CHUNK_METADATA_KERNEL.warmup(vllm_config)


-def _warm_combine_topk_swa_indices_kernel(
-    device: torch.device,
-    num_tokens: int,
-    window_size: int,
-    compress_ratio: int,
-    topk: int,
-    topk_width: int,
-    n: int,
+def _compile_combine_topk_swa_indices_kernel(
+    vllm_config: "VllmConfig",
 ) -> None:
-    from vllm.models.deepseek_v4.common.ops.cache_utils import combine_topk_swa_indices
-
-    if num_tokens <= 0:
-        return
-
-    def _make_topk_indices(*, offset: bool) -> torch.Tensor:
-        if offset:
-            topk_storage = torch.full(
-                (num_tokens * topk_width + 1,),
-                -1,
-                dtype=torch.int32,
-                device=device,
-            )
-            topk_indices = topk_storage[1:].reshape(num_tokens, topk_width)
-        else:
-            topk_indices = torch.full(
-                (num_tokens, topk_width), -1, dtype=torch.int32, device=device
-            )
-        if topk > 0:
-            topk_indices.copy_(
-                torch.arange(num_tokens * topk_width, dtype=torch.int32, device=device)
-                .reshape(num_tokens, topk_width)
-                .remainder(topk_width)
-            )
-        return topk_indices
-
-    query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
-    seq_lens = torch.tensor(
-        [window_size + num_tokens], dtype=torch.int32, device=device
-    )
-    gather_lens = torch.tensor(
-        [min(window_size + num_tokens, window_size + num_tokens - 1)],
-        dtype=torch.int32,
-        device=device,
-    )
-    offset_query_start_loc = torch.empty(3, dtype=torch.int32, device=device)[1:]
-    offset_query_start_loc.copy_(query_start_loc)
-    offset_seq_lens = torch.empty(2, dtype=torch.int32, device=device)[1:]
-    offset_seq_lens.copy_(seq_lens)
-    offset_gather_lens = torch.empty(2, dtype=torch.int32, device=device)[1:]
-    offset_gather_lens.copy_(gather_lens)
-
-    for (
-        offset_topk,
-        offset_query_and_seq,
-        offset_gather,
-    ) in _COMBINE_TOPK_SWA_INPUT_VARIANTS:
-        warmup_topk_indices = _make_topk_indices(offset=offset_topk)
-        warmup_query_start_loc = (
-            offset_query_start_loc if offset_query_and_seq else query_start_loc
-        )
-        warmup_seq_lens = offset_seq_lens if offset_query_and_seq else seq_lens
-        warmup_gather_lens = offset_gather_lens if offset_gather else gather_lens
-        n_values = (n,) if n == 1 else (n, n + 1)
-        for m in (window_size + num_tokens, topk_width):
-            for n_value in n_values:
-                combine_topk_swa_indices(
-                    warmup_topk_indices,
-                    warmup_query_start_loc,
-                    warmup_seq_lens,
-                    warmup_gather_lens,
-                    window_size,
-                    compress_ratio,
-                    topk,
-                    M=m,
-                    N=n_value,
-                )
-
-
-@torch.inference_mode()
-def sparse_mla_triton_warmup(
-    runner: "GPUModelRunner",
-    num_tokens: int,
-    *,
-    compress_ratios: tuple[int, ...],
-    combine_topk_swa_cases: tuple[tuple[int, int, int, int], ...] = (),
-) -> None:
-    device = getattr(runner, "device", torch.device("cuda"))
-    window_size = _hf_config_int(runner, "sliding_window", 128)
-
-    _warm_sparse_swa_prefill_metadata_kernel(device, window_size, num_tokens)
-    for compress_ratio in compress_ratios:
-        _warm_prefill_chunk_metadata_kernel(device, compress_ratio, num_tokens)
-    for compress_ratio, topk, topk_width, n in combine_topk_swa_cases:
-        _warm_combine_topk_swa_indices_kernel(
-            device,
-            num_tokens,
-            window_size,
-            compress_ratio,
-            topk,
-            topk_width,
-            n,
-        )
-
-
-def deepseek_v4_sparse_triton_warmup(
-    runner: "GPUModelRunner",
-    num_tokens: int,
-) -> None:
-    sparse_mla_triton_warmup(
-        runner,
-        num_tokens,
-        compress_ratios=_DSV4_PREFILL_CHUNK_METADATA_COMPRESS_RATIOS,
-        combine_topk_swa_cases=_DSV4_COMBINE_TOPK_SWA_WARMUP_CASES,
+    from vllm.models.deepseek_v4.common.ops.cache_utils import (
+        _COMBINE_TOPK_SWA_INDICES_KERNEL,
    )

+    _COMBINE_TOPK_SWA_INDICES_KERNEL.warmup(vllm_config)

-def sparse_mla_triton_warmup_if_needed(worker: "Worker") -> None:
+
+def sparse_mla_triton_warmup(worker: "Worker") -> None:
    runner = worker.model_runner
    if runner.is_pooling_model:
        return

    max_tokens = worker.scheduler_config.max_num_batched_tokens
-    num_tokens = _clamp_warmup_tokens(8, max_tokens)
-    if num_tokens <= 0:
+    max_num_prefills = min(worker.scheduler_config.max_num_seqs, max_tokens)
+    if max_tokens <= 0 or max_num_prefills <= 0:
        return

+    vllm_config = runner.vllm_config
    try:
        if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS):
-            deepseek_v4_sparse_triton_warmup(runner, num_tokens)
+            _compile_sparse_swa_prefill_metadata_kernel(vllm_config)
+            _compile_prefill_chunk_metadata_kernel(vllm_config)
+            _compile_combine_topk_swa_indices_kernel(vllm_config)
        elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS):
-            sparse_mla_triton_warmup(
-                runner,
-                num_tokens,
-                compress_ratios=(1,),
-            )
+            _compile_sparse_swa_prefill_metadata_kernel(vllm_config)
+            _compile_prefill_chunk_metadata_kernel(vllm_config)
        elif _has_attention_backend(runner, _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS):
-            _warm_prefill_chunk_metadata_kernel(
-                getattr(runner, "device", torch.device("cuda")),
-                compress_ratio=1,
-                query_len=num_tokens,
-            )
+            _compile_prefill_chunk_metadata_kernel(vllm_config)
+
    except Exception:
        logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True)
@@ -14,14 +14,23 @@ preparation.
  window indices for sparse prefill.
 """

+from dataclasses import dataclass
+from typing import Any
+
 import torch

 from vllm.model_executor.layers.quantization.utils.quant_utils import (
    get_fp8_min_max,
 )
+from vllm.model_executor.warmup.jit_warmup import VllmJitKernel, zip_inputs
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    TritonPointerInputVariant,
+    TritonWarmupTensor,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.import_utils import has_cutedsl
+from vllm.utils.math_utils import next_power_of_2


@triton.jit
@@ -528,7 +537,6 @@ def combine_topk_swa_indices(
    N: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    num_tokens = topk_indices.shape[0]
-    num_reqs = seq_lens.shape[0]
    combined_topk = (
        (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
        // _SPARSE_PREFILL_TOPK_ALIGNMENT
@@ -544,13 +552,10 @@ def combine_topk_swa_indices(
        num_tokens, dtype=torch.int32, device=topk_indices.device
    )

-    NUM_WORKERS = 128
-    _combine_topk_swa_indices_kernel[(num_reqs, NUM_WORKERS)](
+    _COMBINE_TOPK_SWA_INDICES_KERNEL(
        combined_indices,
-        combined_indices.stride(0),
        combined_lens,
        topk_indices,
-        topk_indices.stride(0),
        query_start_loc,
        seq_lens,
        gather_lens,
@@ -559,82 +564,245 @@ def combine_topk_swa_indices(
        TOP_K=topk,
        COMPRESS_RATIO=compress_ratio,
        WINDOW_SIZE=window_size,
-        PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
    )
    return combined_indices, combined_lens


-@triton.jit
-def _combine_topk_swa_indices_kernel(
-    combined_indices_ptr,
-    combined_indices_stride,
-    combined_lens_ptr,
-    topk_indices_ptr,
-    topk_indices_stride,
-    query_start_loc_ptr,
-    seq_lens_ptr,
-    gather_lens_ptr,
-    M,
-    N,
-    TOP_K: tl.constexpr,
-    COMPRESS_RATIO: tl.constexpr,
-    WINDOW_SIZE: tl.constexpr,
-    PADDED_TOP_K: tl.constexpr,
+_COMBINE_TOPK_SWA_NUM_WORKERS = 128
+
+
+# Representative pointer alignment variants for Triton pointer specialization.
+_COMBINE_TOPK_SWA_POINTER_INPUTS = zip_inputs(
+    dict(
+        topk_indices=True,
+        query_start_loc=True,
+        seq_lens=True,
+        gather_lens=True,
+    ),
+    dict(
+        topk_indices=True,
+        query_start_loc=False,
+        seq_lens=False,
+        gather_lens=True,
+    ),
+    dict(
+        topk_indices=False,
+        query_start_loc=False,
+        seq_lens=False,
+        gather_lens=False,
+    ),
+)
+
+
+_DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS = zip_inputs(
+    # DSv4-Flash / SWA-only and C4A.
+    dict(compress_ratio=1, topk=0, topk_width=512),
+    dict(compress_ratio=4, topk=512, topk_width=512),
+    # DSv4-Pro C4A.
+    dict(compress_ratio=4, topk=1024, topk_width=1024),
+    # DSv4-Pro C128A.
+    dict(compress_ratio=128, topk=8192, topk_width=8192),
+)
+
+
+def _hf_config_int(vllm_config: Any, name: str, default: int) -> int:
+    model_config = getattr(vllm_config, "model_config", None)
+    hf_config = getattr(model_config, "hf_config", None)
+    return int(getattr(hf_config, name, default) or default)
+
+
+def _scheduler_config_int(vllm_config: Any, name: str, default: int) -> int:
+    scheduler_config = getattr(vllm_config, "scheduler_config", None)
+    return int(getattr(scheduler_config, name, default) or default)
+
+
+class CombineTopkSwaIndicesKernel(
+    VllmJitKernel["CombineTopkSwaIndicesKernel.CompileKey"]
 ):
-    batch_idx = tl.program_id(0)
-    worker_id = tl.program_id(1)
-    num_workers = tl.num_programs(1)
+    @dataclass(frozen=True)
+    class CompileKey:
+        TOP_K: int
+        COMPRESS_RATIO: int
+        WINDOW_SIZE: int
+        PADDED_TOP_K: int
+        input_variant: TritonPointerInputVariant

-    # query_start_loc is a global tensor; rebase to chunk-local offsets
-    # by subtracting the chunk's starting value.
-    base = tl.load(query_start_loc_ptr)
-    query_start = tl.load(query_start_loc_ptr + batch_idx) - base
-    query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
-    query_len = query_end - query_start
-    seq_len = tl.load(seq_lens_ptr + batch_idx)
-    gather_len = tl.load(gather_lens_ptr + batch_idx)
-    start_pos = seq_len - query_len
-    # The SWA portion of the gathered buffer starts from position
-    # (seq_len - gather_len), not position 0. We need this offset
-    # to correctly index into the gathered buffer.
-    gather_start = seq_len - gather_len
+    @staticmethod
+    @triton.jit(
+        do_not_specialize=[
+            "combined_indices_stride",
+            "topk_indices_stride",
+            "M",
+            "N",
+        ]
+    )
+    def kernel(
+        combined_indices_ptr,
+        combined_indices_stride,
+        combined_lens_ptr,
+        topk_indices_ptr,
+        topk_indices_stride,
+        query_start_loc_ptr,
+        seq_lens_ptr,
+        gather_lens_ptr,
+        M,
+        N,
+        TOP_K: tl.constexpr,
+        COMPRESS_RATIO: tl.constexpr,
+        WINDOW_SIZE: tl.constexpr,
+        PADDED_TOP_K: tl.constexpr,
+    ):
+        batch_idx = tl.program_id(0)
+        worker_id = tl.program_id(1)
+        num_workers = tl.num_programs(1)

-    for token_idx in range(query_start + worker_id, query_end, num_workers):
-        # topk_len is fully determined by the query token's absolute position:
-        # both the C4A indexer and the C128A metadata builder emit
-        # min((pos + 1) // compress_ratio, topk_tokens) valid entries.
-        # Caller passes TOP_K=0 for SWA-only layers to zero this out.
-        token_idx_in_query = token_idx - query_start
-        pos = start_pos + token_idx_in_query
-        topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
-        swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
+        # query_start_loc is a global tensor; rebase to chunk-local offsets
+        # by subtracting the chunk's starting value.
+        base = tl.load(query_start_loc_ptr)
+        query_start = tl.load(query_start_loc_ptr + batch_idx) - base
+        query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
+        query_len = query_end - query_start
+        seq_len = tl.load(seq_lens_ptr + batch_idx)
+        gather_len = tl.load(gather_lens_ptr + batch_idx)
+        start_pos = seq_len - query_len
+        # The SWA portion of the gathered buffer starts from position
+        # (seq_len - gather_len), not position 0. We need this offset
+        # to correctly index into the gathered buffer.
+        gather_start = seq_len - gather_len

-        offset = tl.arange(0, PADDED_TOP_K)
-        mask = offset < topk_len
-        topk_indices = tl.load(
-            topk_indices_ptr + token_idx * topk_indices_stride + offset,
-            mask=mask,
+        for token_idx in range(query_start + worker_id, query_end, num_workers):
+            # topk_len is fully determined by the query token's absolute position:
+            # both the C4A indexer and the C128A metadata builder emit
+            # min((pos + 1) // compress_ratio, topk_tokens) valid entries.
+            # Caller passes TOP_K=0 for SWA-only layers to zero this out.
+            token_idx_in_query = token_idx - query_start
+            pos = start_pos + token_idx_in_query
+            topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
+            swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
+
+            offset = tl.arange(0, PADDED_TOP_K)
+            mask = offset < topk_len
+            topk_indices = tl.load(
+                topk_indices_ptr + token_idx * topk_indices_stride + offset,
+                mask=mask,
+            )
+            tl.store(
+                combined_indices_ptr + token_idx * combined_indices_stride + offset,
+                topk_indices + M * batch_idx,
+                mask=mask,
+            )
+            offset = tl.arange(0, WINDOW_SIZE)
+            # Index into gathered buffer: N + (position - gather_start)
+            # For positions [pos - swa_len + 1, pos], the buffer indices are:
+            # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start]
+            tl.store(
+                combined_indices_ptr
+                + token_idx * combined_indices_stride
+                + topk_len
+                + offset,
+                M * batch_idx + N + offset + pos - swa_len + 1 - gather_start,
+                mask=offset < swa_len,
+            )
+
+            combined_len = topk_len + swa_len
+            tl.store(combined_lens_ptr + token_idx, combined_len)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        topk_width: int,
+        topk_indices: bool,
+        query_start_loc: bool,
+        seq_lens: bool,
+        gather_lens: bool,
+        topk: int,
+        compress_ratio: int,
+        WINDOW_SIZE: int,
+    ) -> CompileKey:
+        padded_topk = next_power_of_2(topk_width)
+        input_variant = TritonPointerInputVariant.from_alignment(
+            topk_indices=topk_indices,
+            query_start_loc=query_start_loc,
+            seq_lens=seq_lens,
+            gather_lens=gather_lens,
        )
-        tl.store(
-            combined_indices_ptr + token_idx * combined_indices_stride + offset,
-            topk_indices + M * batch_idx,
-            mask=mask,
-        )
-        offset = tl.arange(0, WINDOW_SIZE)
-        # Index into gathered buffer: N + (position - gather_start)
-        # For positions [pos - swa_len + 1, pos], the buffer indices are:
-        # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start]
-        tl.store(
-            combined_indices_ptr
-            + token_idx * combined_indices_stride
-            + topk_len
-            + offset,
-            M * batch_idx + N + offset + pos - swa_len + 1 - gather_start,
-            mask=offset < swa_len,
+        return self.CompileKey(
+            TOP_K=topk,
+            COMPRESS_RATIO=compress_ratio,
+            WINDOW_SIZE=WINDOW_SIZE,
+            PADDED_TOP_K=padded_topk,
+            input_variant=input_variant,
        )

-        combined_len = topk_len + swa_len
-        tl.store(combined_lens_ptr + token_idx, combined_len)
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        if _scheduler_config_int(vllm_config, "max_num_batched_tokens", 0) <= 0:
+            return []
+
+        window_size = _hf_config_int(vllm_config, "sliding_window", 128)
+        return self._trace_dispatch(self.dispatch)(
+            _DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS,
+            _COMBINE_TOPK_SWA_POINTER_INPUTS,
+            WINDOW_SIZE=window_size,
+        )
+
+    def compile(self, compile_key: CompileKey) -> None:
+        warmup = getattr(self.kernel, "warmup", None)
+        assert warmup is not None
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        input_variant = compile_key.input_variant
+        warmup(
+            int32_ptr,
+            1,  # do not specialize combined_indices_stride
+            int32_ptr,
+            input_variant.pointer("topk_indices", torch.int32),
+            1,  # do not specialize topk_indices_stride
+            input_variant.pointer("query_start_loc", torch.int32),
+            input_variant.pointer("seq_lens", torch.int32),
+            input_variant.pointer("gather_lens", torch.int32),
+            1,  # do not specialize M
+            1,  # do not specialize N
+            TOP_K=compile_key.TOP_K,
+            COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
+            WINDOW_SIZE=compile_key.WINDOW_SIZE,
+            PADDED_TOP_K=compile_key.PADDED_TOP_K,
+            grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),
+        )
+
+    def __call__(
+        self,
+        combined_indices: torch.Tensor,
+        combined_lens: torch.Tensor,
+        topk_indices: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        gather_lens: torch.Tensor,
+        M: int,
+        N: int,
+        *,
+        TOP_K: int,
+        COMPRESS_RATIO: int,
+        WINDOW_SIZE: int,
+    ) -> None:
+        num_reqs = seq_lens.shape[0]
+        self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](
+            combined_indices,
+            combined_indices.stride(0),
+            combined_lens,
+            topk_indices,
+            topk_indices.stride(0),
+            query_start_loc,
+            seq_lens,
+            gather_lens,
+            M,
+            N,
+            TOP_K=TOP_K,
+            COMPRESS_RATIO=COMPRESS_RATIO,
+            WINDOW_SIZE=WINDOW_SIZE,
+            PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]),
+        )
+
+
+_COMBINE_TOPK_SWA_INDICES_KERNEL = CombineTopkSwaIndicesKernel()


 def build_flashinfer_mixed_sparse_indices(
@@ -476,15 +476,17 @@ class MediaConnector:
        self,
        image_url: str,
        *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
    ) -> Image.Image:
        """
        Load a PIL image from an HTTP or base64 data URL.

-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
        """
        image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
        )

        try:
@@ -501,15 +503,17 @@ class MediaConnector:
        self,
        image_url: str,
        *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
    ) -> Image.Image:
        """
        Asynchronously load a PIL image from an HTTP or base64 data URL.

-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
        """
        image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
        )

        try:
@@ -526,14 +530,14 @@ class MediaConnector:
        self,
        video_url: str,
        *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
        video_processor: str | None = None,
    ) -> tuple[npt.NDArray, dict[str, Any]]:
        """
        Load video from an HTTP or base64 data URL.
        """
        image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
        )
        video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
        if "video_backend" not in video_io_kwargs and (
@@ -552,16 +556,18 @@ class MediaConnector:
        self,
        video_url: str,
        *,
-        image_mode: str = "RGB",
+        image_mode: str | None = "RGB",
        video_processor: str | None = None,
    ) -> tuple[npt.NDArray, dict[str, Any]]:
        """
        Asynchronously load video from an HTTP or base64 data URL.

-        By default, the image is converted into RGB format.
+        By default, the image is converted into RGB format. Set
+        `media_io_kwargs={"image": {"image_mode": None}}` to keep the
+        original image mode (e.g. preserving the alpha channel).
        """
        image_io = ImageMediaIO(
-            image_mode=image_mode, **self.media_io_kwargs.get("image", {})
+            **({"image_mode": image_mode} | self.media_io_kwargs.get("image", {}))
        )
        video_io_kwargs = dict(self.media_io_kwargs.get("video", {}))
        if "video_backend" not in video_io_kwargs and (
@@ -25,9 +25,11 @@ class ImageMediaIO(MediaIO[Image.Image]):
    error handling.
    """

-    def __init__(self, image_mode: str = "RGB", **kwargs) -> None:
+    def __init__(self, image_mode: str | None = "RGB", **kwargs) -> None:
        super().__init__()

+        # Target mode for loaded images; `None` keeps the original mode
+        # (i.e. no conversion, alpha channel is preserved as-is).
        self.image_mode = image_mode
        # `kwargs` contains custom arguments from
        # --media-io-kwargs for this modality, merged with
@@ -62,7 +64,7 @@ class ImageMediaIO(MediaIO[Image.Image]):
        """Convert image mode with custom background color."""
        if isinstance(image, MediaWithBytes):
            image = image.media
-        if image.mode == self.image_mode:
+        if self.image_mode is None or image.mode == self.image_mode:
            return image
        elif image.mode == "RGBA" and self.image_mode == "RGB":
            return rgba_to_rgb(image, self.rgba_background_color)
@@ -58,13 +58,14 @@ def encode_audio_url(
 def encode_image_base64(
    image: Image.Image,
    *,
-    image_mode: str = "RGB",
+    image_mode: str | None = "RGB",
    format: str = "PNG",
 ) -> str:
    """
    Encode a pillow image to base64 format.

    By default, the image is converted into RGB format before being encoded.
+    Pass `image_mode=None` to keep the original image mode.
    """
    image_io = ImageMediaIO(image_mode=image_mode)
    return image_io.encode_base64(image, image_format=format)
@@ -73,13 +74,14 @@ def encode_image_base64(
 def encode_image_url(
    image: Image.Image,
    *,
-    image_mode: str = "RGB",
+    image_mode: str | None = "RGB",
    format: str = "PNG",
 ) -> str:
    """
    Encode a pillow image as a data URL.

    By default, the image is converted into RGB format before being encoded.
+    Pass `image_mode=None` to keep the original image mode.
    """
    image_b64 = encode_image_base64(image, image_mode=image_mode, format=format)
    mimetype = mimetypes.types_map.get("." + format.lower(), "image")
@@ -14,6 +14,7 @@ from datetime import timedelta
 from functools import cache, lru_cache, wraps
 from typing import TYPE_CHECKING, NamedTuple, TypeVar

+import regex as re
 import torch
 from torch.distributed import PrefixStore, ProcessGroup
 from torch.distributed.distributed_c10d import is_nccl_available
@@ -48,6 +49,8 @@ _R = TypeVar("_R")

 pynvml = import_pynvml()

+_COMPILED_CUDA_ARCH_RE = re.compile(r"^(?P<major>\d+)\.(?P<minor>\d+)(?P<kind>[af])?$")
+
 # pytorch 2.5 uses cudnn sdpa by default, which will cause crash on some models
 # see https://github.com/huggingface/diffusers/issues/9704 for details
 torch.backends.cuda.enable_cudnn_sdp(False)
@@ -282,6 +285,62 @@ class CudaPlatformBase(Platform):
    def log_warnings(cls):
        pass

+    @staticmethod
+    def _compiled_arch_covers_device(arch: str, capability: DeviceCapability) -> bool:
+        match = _COMPILED_CUDA_ARCH_RE.match(arch)
+        if match is None:
+            return False
+
+        major = int(match.group("major"))
+        minor = int(match.group("minor"))
+        kind = match.group("kind")
+
+        if kind == "f":
+            return major == capability.major
+        return major == capability.major and minor == capability.minor
+
+    @classmethod
+    def _warn_if_device_arch_not_compiled(cls) -> None:
+        try:
+            compiled_archs_str = torch.ops._C.get_compiled_cuda_archs()
+        except (AttributeError, RuntimeError):
+            return
+
+        compiled_archs = [
+            arch.strip() for arch in compiled_archs_str.split(",") if arch.strip()
+        ]
+        if not compiled_archs:
+            return
+
+        unsupported_devices: list[str] = []
+        for device_id in range(cls.device_count()):
+            capability = cls.get_device_capability(device_id)
+            if capability is None:
+                continue
+            if any(
+                cls._compiled_arch_covers_device(arch, capability)
+                for arch in compiled_archs
+            ):
+                continue
+            device_name = cls.get_device_name(device_id)
+            unsupported_devices.append(
+                f"{device_id}: {device_name} (compute capability "
+                f"{capability.as_version_str()})"
+            )
+
+        if unsupported_devices:
+            logger.warning_once(
+                "The current vLLM wheel was built for CUDA architectures %s, "
+                "but the following visible CUDA device(s) are not covered: %s. "
+                "This may cause CUDA kernel launch failures or force slower "
+                "fallback paths. Note that CUDA 13 family-specific targets "
+                "such as 10.0f and 12.0f cover the corresponding major-version "
+                "GPU family, while architecture-specific targets such as 12.1a "
+                "cover only that exact compute capability.",
+                ", ".join(compiled_archs),
+                "; ".join(unsupported_devices),
+            )
+
    @classmethod
    def is_pin_memory_available(cls) -> bool:
        if in_wsl():
@@ -290,7 +349,8 @@ class CudaPlatformBase(Platform):
            # kernel with limited pinned memory support for CUDA.
            version = _get_wsl_kernel_version()
            if version is None or version < (4, 19, 121):
-                logger.warning_once(
+                # warning_once() causes a circular import on WSL, see #48397.
+                logger.warning(
                    "Using 'pin_memory=False' as WSL is detected and the "
                    "WSL2 kernel version is below 4.19.121. This may slow "
                    "down performance. Please run `wsl --update`."
@@ -948,6 +1008,7 @@ class NvmlCudaPlatform(CudaPlatformBase):
    @classmethod
    @with_nvml_context
    def log_warnings(cls):
+        cls._warn_if_device_arch_not_compiled()
        device_ids: int = pynvml.nvmlDeviceGetCount()
        if device_ids > 1:
            device_names = [cls._get_physical_device_name(i) for i in range(device_ids)]
@@ -991,7 +991,8 @@ class Platform:
            # Pinned memory support under WSL depends on the vendor and driver
            # version. Conservative default: return False. Platform subclasses
            # that can verify support (e.g. CudaPlatformBase) override this.
-            logger.warning_once(
+            # warning_once() causes a circular import on WSL, see #48397.
+            logger.warning(
                "Using 'pin_memory=False' as WSL is detected. "
                "This may slow down performance."
            )
@@ -106,6 +106,9 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
        DFlash drafter. Mapped to both eagle_aux_hidden_state_layer_ids
        (for gpu_model_runner) and dflash_config.target_layer_ids (for the
        DFlash model).
+    - sample_from_anchor: Whether to sample from the anchor position. Default
+        False (anchor is a bonus token, only mask tokens predict, yielding
+        block_size - 1 speculative tokens).
    """
    pre_trained_config["architectures"] = ["DFlashDraftModel"]
    pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
@@ -119,6 +122,7 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
    pre_trained_config["dflash_config"] = {
        "mask_token_id": config_dict["mask_token_id"],
        "target_layer_ids": [i - 1 for i in aux_layer_ids],
+        "sample_from_anchor": config_dict.get("sample_from_anchor", False),
    }
    # Enable causal masking in SWA for vllm-project/speculators models
    pre_trained_config["dflash_config"]["causal"] = not config_dict.get(
@@ -146,10 +150,14 @@ def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
    - aux_hidden_state_layer_ids (required): target layer indices feeding the
        drafter. Mapped to both eagle_aux_hidden_state_layer_ids and
        target_layer_ids (DSpark's i-1 layer semantics).
+    - sample_from_anchor: Whether to sample from the anchor position. Default
+        False (anchor is a bonus token, only mask tokens predict, yielding
+        block_size - 1 speculative tokens).
    """
    pre_trained_config["architectures"] = ["Qwen3DSparkModel"]
-    # Speculators DSpark uses the 1+N fill-in block (anchor is a bonus token).
-    pre_trained_config["dspark_bonus_anchor"] = True
+    pre_trained_config["sample_from_anchor"] = config_dict.get(
+        "sample_from_anchor", False
+    )

    aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
    pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
@@ -110,6 +110,7 @@ class MemorySnapshot:
    """Memory snapshot."""

    torch_peak: int = 0
+    torch_allocated: int = 0
    free_memory: int = 0
    total_memory: int = 0
    cuda_memory: int = 0
@@ -139,9 +140,9 @@ class MemorySnapshot:
        # After `torch.accelerator.reset_peak_memory_stats()`,
        # `torch.accelerator.memory_reserved()` will keep growing, and only shrink
        # when we call `torch.accelerator.empty_cache()` or OOM happens.
-        self.torch_peak = torch.accelerator.memory_stats(device).get(
-            "allocated_bytes.all.peak", 0
-        )
+        stats = torch.accelerator.memory_stats(device)
+        self.torch_peak = stats.get("allocated_bytes.all.peak", 0)
+        self.torch_allocated = stats.get("allocated_bytes.all.current", 0)

        self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device)
        if current_platform.is_integrated_gpu(device.index):
@@ -172,6 +173,7 @@ class MemorySnapshot:

        return MemorySnapshot(
            torch_peak=self.torch_peak - other.torch_peak,
+            torch_allocated=self.torch_allocated - other.torch_allocated,
            free_memory=self.free_memory - other.free_memory,
            total_memory=self.total_memory - other.total_memory,
            cuda_memory=self.cuda_memory - other.cuda_memory,
@@ -185,6 +187,7 @@ class MemorySnapshot:
    def __repr__(self) -> str:
        return (
            f"torch_peak={format_gib(self.torch_peak)}GiB, "
+            f"torch_allocated={format_gib(self.torch_allocated)}GiB, "
            f"free_memory={format_gib(self.free_memory)}GiB, "
            f"total_memory={format_gib(self.total_memory)}GiB, "
            f"{current_platform.device_name}_memory={format_gib(self.cuda_memory)}GiB, "
@@ -202,6 +205,8 @@ class MemoryProfilingResult:
    non_kv_cache_memory: int = 0
    torch_peak_increase: int = 0
    non_torch_increase: int = 0
+    total_consumed: int = 0
+    transient_peak_headroom: int = 0
    weights_memory: int = 0
    before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
    profile_time: float = 0.0
@@ -219,8 +224,8 @@ class MemoryProfilingResult:
            f"{format_gib(self.non_kv_cache_memory)}GiB; "
            f"torch peak memory increase: "
            f"{format_gib(self.torch_peak_increase)}GiB; "
-            f"non-torch forward increase memory: "
-            f"{format_gib(self.non_torch_increase)}GiB; "
+            f"total consumed (from mem_get_info): "
+            f"{format_gib(self.total_consumed)}GiB; "
            f"weights memory: {format_gib(self.weights_memory)}GiB."
        )

@@ -306,8 +311,16 @@ def memory_profiling(
    result.non_torch_increase = diff_from_create.non_torch_memory
    result.profile_time = diff_profile.timestamp

-    non_torch_memory = result.non_torch_increase
-    peak_activation_memory = result.torch_peak_increase
-    result.non_kv_cache_memory = (
-        non_torch_memory + peak_activation_memory + result.weights_memory
+    # Measure total consumption via mem_get_info() instead of
+    # memory_reserved(), which goes negative when pluggable allocators
+    # (e.g. cumem) bypass PyTorch's tracking.
+    result.total_consumed = (
+        result.before_create.free_memory - result.after_profile.free_memory
    )
+
+    # total_consumed already covers persistent torch allocations; add only the
+    # transient peak headroom to avoid double-counting.
+    result.transient_peak_headroom = (
+        result.after_profile.torch_peak - result.after_profile.torch_allocated
+    )
+    result.non_kv_cache_memory = result.total_consumed + result.transient_peak_headroom
@@ -273,6 +273,29 @@ def is_fa_version_supported(fa_version: int) -> bool:
        return False


+def flash_attn_supports_kv_cache_dtype(
+    kv_cache_dtype: str = "fp8_e4m3",
+    *,
+    requires_alibi: bool = False,
+    head_size: int | None = None,
+    head_size_v: int | None = None,
+    has_sinks: bool = False,
+) -> bool:
+    if kv_cache_dtype == "fp8_e5m2":
+        return False
+    if current_platform.is_xpu():
+        return True
+    fa_version = get_flash_attn_version(
+        requires_alibi=requires_alibi,
+        head_size=head_size,
+        head_size_v=head_size_v,
+        has_sinks=has_sinks,
+    )
+    return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (
+        fa_version == 4 and current_platform.is_device_capability_family(100)
+    )
+
+
 def flash_attn_supports_quant_query_input() -> bool:
    return not current_platform.is_xpu()

@@ -22,6 +22,7 @@ from vllm.v1.attention.backend import (
    MultipleOf,
 )
 from vllm.v1.attention.backends.fa_utils import (
+    flash_attn_supports_kv_cache_dtype,
    flash_attn_supports_quant_query_input,
    get_flash_attn_version,
    is_fa_version_supported,
@@ -74,6 +75,8 @@ class FlashAttentionBackend(AttentionBackend):
        "auto",
        "float16",
        "bfloat16",
+        "fp8",
+        "fp8_e4m3",
    ]

    @staticmethod
@@ -178,14 +181,11 @@ class FlashAttentionBackend(AttentionBackend):
    def supports_kv_cache_dtype(cls, kv_cache_dtype: CacheDType | None) -> bool:
        if kv_cache_dtype is None:
            return True
-        if kv_cache_dtype in ("fp8", "fp8_e4m3"):
-            if current_platform.is_xpu():
-                return True
-            return (
-                get_flash_attn_version() == 3
-                and current_platform.is_device_capability_family(90)
-            )
-        return kv_cache_dtype in ["auto", "float16", "bfloat16"]
+        if kv_cache_dtype not in cls.supported_kv_cache_dtypes:
+            return False
+        if is_quantized_kv_cache(kv_cache_dtype):
+            return flash_attn_supports_kv_cache_dtype(kv_cache_dtype)
+        return True

    @classmethod
    def supports_mm_prefix(cls) -> bool:
@@ -216,6 +216,17 @@ class FlashAttentionBackend(AttentionBackend):
    ) -> str | None:
        if has_sink and device_capability < DeviceCapability(9, 0):
            return "sink not supported on compute capability < 9.0"
+        if (
+            kv_cache_dtype is not None
+            and is_quantized_kv_cache(kv_cache_dtype)
+            and not flash_attn_supports_kv_cache_dtype(
+                kv_cache_dtype,
+                head_size=head_size,
+                head_size_v=head_size,
+                has_sinks=has_sink,
+            )
+        ):
+            return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100"
        if (
            use_mm_prefix
            and get_flash_attn_version(head_size=head_size, has_sinks=has_sink) != 4
@@ -772,6 +783,7 @@ class FlashAttentionImpl(AttentionImpl):
        self.vllm_flash_attn_version = get_flash_attn_version(
            requires_alibi=alibi_slopes is not None,
            head_size=head_size,
+            has_sinks=sinks is not None,
        )
        logger.info_once(
            "Using FlashAttention version %s",
@@ -780,6 +792,20 @@ class FlashAttentionImpl(AttentionImpl):
        # Cache the batch invariant result for use in forward passes
        self.batch_invariant_enabled = envs.VLLM_BATCH_INVARIANT

+        if is_quantized_kv_cache(
+            self.kv_cache_dtype
+        ) and not flash_attn_supports_kv_cache_dtype(
+            self.kv_cache_dtype,
+            requires_alibi=alibi_slopes is not None,
+            head_size=head_size,
+            head_size_v=head_size,
+            has_sinks=sinks is not None,
+        ):
+            raise NotImplementedError(
+                f"FlashAttention does not support {self.kv_cache_dtype}"
+                " kv-cache on this device."
+            )
+
        self.sinks = sinks
        if self.sinks is not None:
            assert flash_attn_supports_sinks(), (
@@ -8,6 +8,14 @@ import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
+from vllm.model_executor.warmup.jit_warmup import (
+    VllmJitKernel,
+    WarmupIntRange,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    TritonPointerInputVariant,
+    TritonWarmupTensor,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.deep_gemm import (
@@ -190,6 +198,211 @@ class DeepseekV32IndexerPrefillChunkMetadata:
    max_local_total_seq_lens: int = 0


+_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS = (
+    TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
+    TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
+)
+
+
+class BuildPrefillChunkMetadataKernel(
+    VllmJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
+):
+    BLOCK_SIZE = 1024
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        query_slice_start: int
+        query_slice_stop: int
+        DCP_RANK: int
+        DCP_WORLD: int
+        DCP_INTERLEAVE: int
+        BLOCK_SIZE: int
+        COMPRESS_RATIO: int
+        input_variant: TritonPointerInputVariant
+
+    @staticmethod
+    @triton.jit
+    def kernel(
+        # Inputs
+        query_start_loc_ptr,
+        uncompressed_seq_lens_ptr,
+        cu_compressed_seq_lens_ptr,
+        # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP,
+        # aliases cu_compressed_seq_lens_ptr otherwise.
+        row_start_cu_compressed_seq_lens_ptr,
+        # Outputs
+        token_to_seq_ptr,
+        cu_compressed_seq_len_ks_ptr,
+        cu_compressed_seq_len_ke_ptr,
+        query_slice_start,
+        query_slice_stop,
+        DCP_RANK,
+        DCP_WORLD,
+        DCP_INTERLEAVE,
+        BLOCK_SIZE: tl.constexpr,
+        COMPRESS_RATIO: tl.constexpr,
+    ):
+        batch_idx = tl.program_id(0)
+
+        query_start = tl.load(query_start_loc_ptr + batch_idx)
+        query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
+        query_len = query_end - query_start
+
+        seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
+        seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
+        compressed_seq_len = seq_end - seq_start
+
+        # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start
+        # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr).
+        row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx)
+
+        uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
+        start_pos = uncompressed_seq_len - query_len
+
+        for i in range(0, query_len, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            abs_pos = query_start + offset
+            mask = (
+                (offset < query_len)
+                & (abs_pos >= query_slice_start)
+                & (abs_pos < query_slice_stop)
+            )
+            out_pos = abs_pos - query_slice_start
+
+            # cu_seq_len_ks: row start in the gathered K buffer.
+            tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask)
+
+            # cu_seq_len_ke: row start + per-token context length. Under DCP the
+            # global per-token length is sharded across ranks.
+            global_ctx = start_pos + 1 + offset
+            len_per_token = global_ctx // COMPRESS_RATIO
+            if DCP_WORLD > 1:
+                # Per-rank local context length under interleave-aware DCP, matching
+                # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
+                base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE
+                remainder = len_per_token - base * DCP_WORLD
+                remainder = tl.minimum(
+                    tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE
+                )
+                len_per_token = base + remainder
+            tl.store(
+                cu_compressed_seq_len_ke_ptr + out_pos,
+                row_start + len_per_token,
+                mask=mask,
+            )
+
+        # Compute token_to_seq
+        for i in range(0, compressed_seq_len, BLOCK_SIZE):
+            offset = i + tl.arange(0, BLOCK_SIZE)
+            mask = offset < compressed_seq_len
+            tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        query_slice_start: int,
+        query_slice_stop: int,
+        DCP_RANK: int,
+        DCP_WORLD: int,
+        DCP_INTERLEAVE: int,
+        BLOCK_SIZE: int,
+        COMPRESS_RATIO: int,
+        input_variant: TritonPointerInputVariant,
+    ) -> CompileKey:
+        return self.CompileKey(
+            query_slice_start=query_slice_start,
+            query_slice_stop=query_slice_stop,
+            DCP_RANK=DCP_RANK,
+            DCP_WORLD=DCP_WORLD,
+            DCP_INTERLEAVE=DCP_INTERLEAVE,
+            BLOCK_SIZE=BLOCK_SIZE,
+            COMPRESS_RATIO=COMPRESS_RATIO,
+            input_variant=input_variant,
+        )
+
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
+        max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8))
+        hf_config = vllm_config.model_config.hf_config
+        parallel_config = vllm_config.parallel_config
+        dcp_world = parallel_config.decode_context_parallel_size
+        dcp_interleave = parallel_config.cp_kv_cache_interleave_size
+        dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
+        compress_ratios = tuple(
+            int(ratio)
+            for ratio in (getattr(hf_config, "compress_ratios", None) or (1,))
+        )
+        return self._trace_dispatch(self.dispatch)(
+            query_slice_start=WarmupIntRange(0, 2),
+            query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
+            DCP_RANK=dcp_rank,
+            DCP_WORLD=dcp_world,
+            DCP_INTERLEAVE=dcp_interleave,
+            BLOCK_SIZE=self.BLOCK_SIZE,
+            COMPRESS_RATIO=list(compress_ratios),
+            input_variant=_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS,
+        )
+
+    def compile(self, compile_key: CompileKey) -> None:
+        warmup = getattr(self.kernel, "warmup", None)
+        assert warmup is not None
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        warmup(
+            int32_ptr,
+            compile_key.input_variant.pointer("uncompressed_seq_lens", torch.int32),
+            int32_ptr,
+            int32_ptr,
+            int32_ptr,
+            int32_ptr,
+            int32_ptr,
+            compile_key.query_slice_start,
+            compile_key.query_slice_stop,
+            compile_key.DCP_RANK,
+            compile_key.DCP_WORLD,
+            compile_key.DCP_INTERLEAVE,
+            BLOCK_SIZE=compile_key.BLOCK_SIZE,
+            COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
+            grid=(1,),
+        )
+
+    def __call__(
+        self,
+        query_start_loc: torch.Tensor,
+        uncompressed_seq_lens: torch.Tensor,
+        cu_compressed_seq_lens: torch.Tensor,
+        row_start_cu_compressed_seq_lens: torch.Tensor,
+        token_to_seq: torch.Tensor,
+        cu_compressed_seq_len_ks: torch.Tensor,
+        cu_compressed_seq_len_ke: torch.Tensor,
+        query_slice_start: int,
+        query_slice_stop: int,
+        DCP_RANK: int,
+        DCP_WORLD: int,
+        DCP_INTERLEAVE: int,
+        *,
+        num_reqs: int,
+        COMPRESS_RATIO: int,
+    ) -> None:
+        self.kernel[(num_reqs,)](
+            query_start_loc,
+            uncompressed_seq_lens,
+            cu_compressed_seq_lens,
+            row_start_cu_compressed_seq_lens,
+            token_to_seq,
+            cu_compressed_seq_len_ks,
+            cu_compressed_seq_len_ke,
+            query_slice_start,
+            query_slice_stop,
+            DCP_RANK,
+            DCP_WORLD,
+            DCP_INTERLEAVE,
+            BLOCK_SIZE=self.BLOCK_SIZE,
+            COMPRESS_RATIO=COMPRESS_RATIO,
+        )
+
+
+_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
+
+
@dataclass
 class DeepseekV32IndexerPrefillMetadata:
    chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
@@ -841,7 +1054,7 @@ def build_prefill_chunk_metadata(

    # Under DCP the kernel writes this rank's local row bounds into
    # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens.
-    _build_prefill_chunk_metadata_kernel[(num_reqs,)](
+    _BUILD_PREFILL_CHUNK_METADATA_KERNEL(
        query_start_loc,
        uncompressed_seq_lens[start_idx:end_idx],
        cu_seq_lens,
@@ -854,7 +1067,7 @@ def build_prefill_chunk_metadata(
        dcp_rank,
        dcp_world_size,
        cp_kv_cache_interleave_size,
-        BLOCK_SIZE=1024,
+        num_reqs=num_reqs,
        COMPRESS_RATIO=compress_ratio,
    )

@@ -881,80 +1094,3 @@ def build_prefill_chunk_metadata(
        local_total_seq_lens=local_total_seq_lens,
        max_local_total_seq_lens=max_local_total_seq_lens,
    )
-
-
-@triton.jit
-def _build_prefill_chunk_metadata_kernel(
-    # Inputs
-    query_start_loc_ptr,
-    uncompressed_seq_lens_ptr,
-    cu_compressed_seq_lens_ptr,
-    # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP,
-    # aliases cu_compressed_seq_lens_ptr otherwise.
-    row_start_cu_compressed_seq_lens_ptr,
-    # Outputs
-    token_to_seq_ptr,
-    cu_compressed_seq_len_ks_ptr,
-    cu_compressed_seq_len_ke_ptr,
-    query_slice_start,
-    query_slice_stop,
-    DCP_RANK,
-    DCP_WORLD,
-    DCP_INTERLEAVE,
-    BLOCK_SIZE: tl.constexpr,
-    COMPRESS_RATIO: tl.constexpr,
-):
-    batch_idx = tl.program_id(0)
-
-    query_start = tl.load(query_start_loc_ptr + batch_idx)
-    query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
-    query_len = query_end - query_start
-
-    seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
-    seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
-    compressed_seq_len = seq_end - seq_start
-
-    # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start
-    # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr).
-    row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx)
-
-    uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
-    start_pos = uncompressed_seq_len - query_len
-
-    for i in range(0, query_len, BLOCK_SIZE):
-        offset = i + tl.arange(0, BLOCK_SIZE)
-        abs_pos = query_start + offset
-        mask = (
-            (offset < query_len)
-            & (abs_pos >= query_slice_start)
-            & (abs_pos < query_slice_stop)
-        )
-        out_pos = abs_pos - query_slice_start
-
-        # cu_seq_len_ks: row start in the gathered K buffer.
-        tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask)
-
-        # cu_seq_len_ke: row start + per-token context length. Under DCP the
-        # global per-token length is sharded across ranks.
-        global_ctx = start_pos + 1 + offset
-        len_per_token = global_ctx // COMPRESS_RATIO
-        if DCP_WORLD > 1:
-            # Per-rank local context length under interleave-aware DCP, matching
-            # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world.
-            base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE
-            remainder = len_per_token - base * DCP_WORLD
-            remainder = tl.minimum(
-                tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE
-            )
-            len_per_token = base + remainder
-        tl.store(
-            cu_compressed_seq_len_ke_ptr + out_pos,
-            row_start + len_per_token,
-            mask=mask,
-        )
-
-    # Compute token_to_seq
-    for i in range(0, compressed_seq_len, BLOCK_SIZE):
-        offset = i + tl.arange(0, BLOCK_SIZE)
-        mask = offset < compressed_seq_len
-        tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)
@@ -3,7 +3,9 @@
 """FlashAttention backend for MLA prefill."""

 import functools
-from typing import TYPE_CHECKING
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any

 import torch

@@ -11,14 +13,7 @@ import vllm.envs as envs
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
    kFp8StaticTensorSym,
 )
-from vllm.model_executor.warmup.cutedsl_warmup import (
-    CuTeDSLCompileUnit,
-    register_cutedsl_warmup_provider,
-)
-from vllm.model_executor.warmup.fa4_cutedsl_config import (
-    FA4MLAPrefillCompileContext,
-    iter_fa4_mla_prefill_compile_requests,
-)
+from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.fa_utils import (
    compile_flash_attn_varlen_func_from_specs,
@@ -32,6 +27,7 @@ from vllm.v1.attention.backends.mla.prefill.base import (

 if TYPE_CHECKING:
    from vllm.config import VllmConfig
+    from vllm.model_executor.layers.attention.mla_attention import MLADims
    from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

 if is_flash_attn_varlen_func_available():
@@ -40,6 +36,245 @@ else:
    flash_attn_varlen_func = None  # type: ignore[assignment]


+FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16)
+
+# Current vLLM MLA prefill expands K/V to num_heads before FA4, so this plan
+# covers qhead_per_kvhead=1. Batch is not a current FA4 MLA-prefill key field.
+# Use b1 for compile-only specs because it is the conservative case for
+# Split-KV shape heuristics.
+# TODO(roberto): FA4 also has direct-GQA and qv/top-k absorbed-MLA paths, but
+# vLLM does not use them in this backend yet; they need a separate
+# num_kv_heads/qv/top-k-aware warmup plan if wired in later.
+FA4_MLA_PREFILL_COMPILE_BATCH_SIZE = 1
+FA4_MLA_PREFILL_Q_TILE = 128
+FA4_MLA_PREFILL_K_TILE = 128
+FA4_MLA_PREFILL_LONG_K_BLOCKS = 32
+FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS = 64
+FA4_MLA_PREFILL_CAUSAL_OPTIONS = (False, True)
+FA4_MLA_PREFILL_LSE_OPTIONS = (False, True)
+
+
+@dataclass(frozen=True)
+class _FA4MLAPrefillShapeProbe:
+    max_seqlen_q: int
+    max_seqlen_k: int
+
+
+class FA4MLAPrefillKernel(VllmJitKernel["FA4MLAPrefillKernel.CompileKey"]):
+    """FA4 MLA prefill compile-key wrapper used by generic JIT warmup."""
+
+    @dataclass(frozen=True)
+    class CompileKey:
+        """High-level FA4 compile-only key used by vLLM warmup.
+
+        This is not the CuTeDSL cache key. FA4 owns the selector that maps these
+        serving inputs to the actual compile-static fields: tile sizes, q_stage,
+        Split-KV, scheduler choice, layout-presence booleans, dtype/head dims,
+        arch, and related fields.
+        """
+
+        q_shape: tuple[int, ...]
+        k_shape: tuple[int, ...]
+        v_shape: tuple[int, ...]
+        q_dtype: torch.dtype
+        max_seqlen_q: int
+        max_seqlen_k: int
+        softmax_scale: float
+        causal: bool
+        fa_version: int
+        v_stride: tuple[int, ...] | None = None
+        cu_seqlens_q_shape: tuple[int, ...] | None = None
+        cu_seqlens_k_shape: tuple[int, ...] | None = None
+        window_size: tuple[int, int] | None = None
+        return_softmax_lse: bool = False
+        num_splits: int = 0
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        batch_size: int,
+        dtype: torch.dtype,
+        num_heads: int,
+        mla_dims: "MLADims",
+        shape_probe: _FA4MLAPrefillShapeProbe,
+        requires_v_padding: bool,
+        causal: bool,
+        fa_version: int,
+        window_size: tuple[int, int] | None = None,
+        return_lse: bool = False,
+        num_splits: int = 0,
+    ) -> CompileKey:
+        return self.CompileKey(
+            q_shape=(
+                batch_size * shape_probe.max_seqlen_q,
+                num_heads,
+                mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
+            ),
+            k_shape=(
+                batch_size * shape_probe.max_seqlen_k,
+                num_heads,
+                mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim,
+            ),
+            v_shape=(
+                batch_size * shape_probe.max_seqlen_k,
+                num_heads,
+                (
+                    mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
+                    if requires_v_padding
+                    else mla_dims.v_head_dim
+                ),
+            ),
+            q_dtype=dtype,
+            max_seqlen_q=shape_probe.max_seqlen_q,
+            max_seqlen_k=shape_probe.max_seqlen_k,
+            softmax_scale=(mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim)
+            ** -0.5,
+            causal=causal,
+            fa_version=fa_version,
+            v_stride=(
+                None
+                if requires_v_padding
+                else (
+                    num_heads * (mla_dims.qk_nope_head_dim + mla_dims.v_head_dim),
+                    mla_dims.qk_nope_head_dim + mla_dims.v_head_dim,
+                    1,
+                )
+            ),
+            cu_seqlens_q_shape=(batch_size + 1,),
+            cu_seqlens_k_shape=(batch_size + 1,),
+            window_size=window_size,
+            return_softmax_lse=return_lse,
+            num_splits=num_splits,
+        )
+
+    def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]:
+        from vllm.model_executor.layers.attention.mla_attention import (
+            get_mla_dims,
+        )
+
+        mla_dims = get_mla_dims(vllm_config.model_config)
+        qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim
+        fa_version = get_flash_attn_version(head_size=qk_head_dim)
+        if fa_version != 4:
+            return []
+
+        dtype = vllm_config.model_config.dtype
+        if dtype not in FA4_STANDARD_DTYPES:
+            dtype = torch.bfloat16
+
+        is_sm90 = current_platform.is_device_capability(90)
+        is_sm100_family = current_platform.is_device_capability_family(100)
+        is_sm120 = current_platform.is_device_capability_family(120)
+        if not (is_sm90 or is_sm100_family or is_sm120):
+            return []
+
+        num_splits = 1 if envs.VLLM_BATCH_INVARIANT else 0
+        if is_sm120 and num_splits != 1:
+            return []
+
+        num_heads = vllm_config.model_config.get_num_attention_heads(
+            vllm_config.parallel_config
+        )
+        if num_heads <= 0:
+            return []
+
+        requires_v_padding = False
+        effective_v_head_dim = (
+            qk_head_dim if requires_v_padding else mla_dims.v_head_dim
+        )
+        shape_probes = [
+            _FA4MLAPrefillShapeProbe(1, FA4_MLA_PREFILL_K_TILE),
+            _FA4MLAPrefillShapeProbe(
+                FA4_MLA_PREFILL_Q_TILE + 1,
+                4 * FA4_MLA_PREFILL_K_TILE,
+            ),
+        ]
+        if num_splits != 1:
+            shape_probes.extend(
+                (
+                    _FA4MLAPrefillShapeProbe(
+                        1,
+                        FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
+                    ),
+                    _FA4MLAPrefillShapeProbe(
+                        FA4_MLA_PREFILL_Q_TILE + 1,
+                        FA4_MLA_PREFILL_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
+                    ),
+                )
+            )
+            if not is_sm90 and qk_head_dim != effective_v_head_dim:
+                shape_probes.extend(
+                    (
+                        _FA4MLAPrefillShapeProbe(
+                            1,
+                            FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
+                        ),
+                        _FA4MLAPrefillShapeProbe(
+                            FA4_MLA_PREFILL_Q_TILE + 1,
+                            FA4_MLA_PREFILL_VERY_LONG_K_BLOCKS * FA4_MLA_PREFILL_K_TILE,
+                        ),
+                    )
+                )
+
+        return self._trace_dispatch(self.dispatch)(
+            batch_size=FA4_MLA_PREFILL_COMPILE_BATCH_SIZE,
+            dtype=dtype,
+            num_heads=num_heads,
+            mla_dims=mla_dims,
+            shape_probe=shape_probes,
+            requires_v_padding=requires_v_padding,
+            causal=FA4_MLA_PREFILL_CAUSAL_OPTIONS,
+            return_lse=FA4_MLA_PREFILL_LSE_OPTIONS,
+            num_splits=num_splits,
+            fa_version=fa_version,
+        )
+
+    @staticmethod
+    def kernel(*args: Any, **kwargs: Any) -> Any:
+        assert flash_attn_varlen_func is not None
+        return flash_attn_varlen_func(*args, **kwargs)
+
+    def compile(self, compile_key: CompileKey) -> None:
+        assert compile_flash_attn_varlen_func_from_specs is not None
+        window_size = (
+            list(compile_key.window_size)
+            if compile_key.window_size is not None
+            else None
+        )
+        compile_flash_attn_varlen_func_from_specs(
+            q_shape=compile_key.q_shape,
+            k_shape=compile_key.k_shape,
+            v_shape=compile_key.v_shape,
+            q_dtype=compile_key.q_dtype,
+            v_stride=compile_key.v_stride,
+            cu_seqlens_q_shape=compile_key.cu_seqlens_q_shape,
+            cu_seqlens_k_shape=compile_key.cu_seqlens_k_shape,
+            max_seqlen_q=compile_key.max_seqlen_q,
+            max_seqlen_k=compile_key.max_seqlen_k,
+            softmax_scale=compile_key.softmax_scale,
+            causal=compile_key.causal,
+            window_size=window_size,
+            return_softmax_lse=compile_key.return_softmax_lse,
+            fa_version=compile_key.fa_version,
+            num_splits=compile_key.num_splits,
+        )
+
+    def __call__(
+        self,
+        *,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        runtime_kernel: Callable[..., Any] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        kernel = self.kernel if runtime_kernel is None else runtime_kernel
+        return kernel(q=q, k=k, v=v, **kwargs)
+
+
+FA4_MLA_PREFILL_KERNEL = FA4MLAPrefillKernel()
+
+
 class FlashAttnPrefillBackend(MLAPrefillBackend):
    """FlashAttention backend for MLA prefill."""

@@ -127,46 +362,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):

        # Track whether we're using vllm's FA or upstream (for ROCm)
        self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu()
-        if self.vllm_flash_attn_version == 4:
-            register_cutedsl_warmup_provider(self)
-
-    def get_cutedsl_warmup_compile_units(self) -> tuple[CuTeDSLCompileUnit, ...]:
-        if self.vllm_flash_attn_version != 4:
-            return ()
-        if compile_flash_attn_varlen_func_from_specs is None:
-            raise RuntimeError(
-                "FA4 compile-only API is unavailable; CuTeDSL warmup does not "
-                "fall back to synthetic forward passes."
-            )
-
-        dtype = self.vllm_config.model_config.dtype
-        if dtype not in self.supported_dtypes:
-            dtype = torch.bfloat16
-
-        qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
-        ctx = FA4MLAPrefillCompileContext(
-            dtype=dtype,
-            num_heads=self.num_heads,
-            qk_head_dim=qk_head_dim,
-            v_head_dim=self.v_head_dim,
-            kv_nope_head_dim=self.qk_nope_head_dim + self.v_head_dim,
-            requires_v_padding=self.requires_v_padding,
-            scale=self.scale,
-            num_splits=1 if envs.VLLM_BATCH_INVARIANT else 0,
-            fa_version=self.vllm_flash_attn_version,
-        )
-        compile_requests = tuple(iter_fa4_mla_prefill_compile_requests(ctx))
-        if not compile_requests:
-            return ()
-
-        return tuple(
-            CuTeDSLCompileUnit(
-                name="fa4_mla_prefill",
-                key=request.key,
-                compile=request.compile,
-            )
-            for request in compile_requests
-        )

    def supports_quant_output(self, quant_key: "QuantKey") -> bool:
        device_capability = current_platform.get_device_capability()
@@ -207,10 +402,11 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
        if envs.VLLM_BATCH_INVARIANT:
            kwargs["num_splits"] = 1

-        attn_out = self.flash_attn_varlen_func(
+        attn_out = FA4_MLA_PREFILL_KERNEL(
            q=q,
            k=k,
            v=maybe_padded_v,
+            runtime_kernel=self.flash_attn_varlen_func,
            softmax_scale=softmax_scale,
            **kwargs,
        )
@@ -7,9 +7,14 @@ import torch

 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
+from vllm.model_executor.warmup.jit_warmup import (
+    VllmJitKernel,
+    WarmupIntRange,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import TritonWarmupTensor
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, next_power_of_2
 from vllm.v1.attention.backend import (
    AttentionBackend,
    AttentionCGSupport,
@@ -279,6 +284,107 @@ class DeepseekSparseSWAMetadata:
        return chunk_plan


+class ComputePrefillMetadataKernel(
+    VllmJitKernel["ComputePrefillMetadataKernel.CompileKey"]
+):
+    @dataclass(frozen=True)
+    class CompileKey:
+        BLOCK_SIZE: int
+
+    @staticmethod
+    @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])
+    def kernel(
+        # Outputs
+        prefill_gather_lens_ptr,
+        # Inputs
+        seq_lens_ptr,
+        query_start_loc_ptr,
+        num_prefills,
+        num_decodes,
+        window_size,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        """Compute prefill gather_lens in a single pass."""
+        offset = tl.arange(0, BLOCK_SIZE)
+        mask = offset < num_prefills
+        # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
+        # masked-off lanes even though the load mask gates the actual read, so
+        # clamp the offset. Caller guarantees num_prefills > 0.
+        safe_offset = tl.minimum(offset, num_prefills - 1)
+
+        seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
+        qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
+        qsl_end = tl.load(
+            query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask
+        )
+
+        query_len = qsl_end - qsl_start
+        prefix_len = seq_len - query_len
+        gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
+
+        tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        num_prefills: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            BLOCK_SIZE=next_power_of_2(num_prefills),
+        )
+
+    def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
+        scheduler_config = vllm_config.scheduler_config
+        max_prefills = max(
+            1,
+            min(
+                scheduler_config.max_num_seqs,
+                scheduler_config.max_num_batched_tokens,
+            ),
+        )
+        return self._trace_dispatch(self.dispatch)(
+            num_prefills=WarmupIntRange(1, max_prefills + 1),
+        )
+
+    def compile(self, compile_key: CompileKey) -> None:
+        warmup = getattr(self.kernel, "warmup", None)
+        assert warmup is not None
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        warmup(
+            int32_ptr,
+            int32_ptr,
+            int32_ptr,
+            compile_key.BLOCK_SIZE,
+            0,
+            1,
+            BLOCK_SIZE=compile_key.BLOCK_SIZE,
+            grid=(1,),
+        )
+
+    def __call__(
+        self,
+        prefill_gather_lens: torch.Tensor,
+        seq_lens: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        num_prefills: int,
+        num_decodes: int,
+        window_size: int,
+    ) -> None:
+        compile_key = self.dispatch(num_prefills=num_prefills)
+        self.kernel[(1,)](
+            prefill_gather_lens,
+            seq_lens,
+            query_start_loc,
+            num_prefills,
+            num_decodes,
+            window_size,
+            BLOCK_SIZE=compile_key.BLOCK_SIZE,
+        )
+
+
+_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
+
+
 class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
    """Builds metadata for DeepseekV4 SWA cache.

@@ -602,14 +708,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
            pfx_gather_lens = torch.empty(
                num_prefills, dtype=torch.int32, device=seq_lens.device
            )
-            _compute_prefill_metadata_kernel[(1,)](
+            _COMPUTE_PREFILL_METADATA_KERNEL(
                pfx_gather_lens,
                seq_lens,
                query_start_loc,
                num_prefills,
                num_decodes,
                self.window_size,
-                BLOCK_SIZE=triton.next_power_of_2(num_prefills),
            )

            result["prefill_seq_lens"] = seq_lens[num_decodes:]
@@ -626,37 +731,6 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
        return result


-@triton.jit
-def _compute_prefill_metadata_kernel(
-    # Outputs
-    prefill_gather_lens_ptr,
-    # Inputs
-    seq_lens_ptr,
-    query_start_loc_ptr,
-    num_prefills,
-    num_decodes,
-    window_size,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """Compute prefill gather_lens in a single pass."""
-    offset = tl.arange(0, BLOCK_SIZE)
-    mask = offset < num_prefills
-    # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for
-    # masked-off lanes even though the load mask gates the actual read, so
-    # clamp the offset. Caller guarantees num_prefills > 0.
-    safe_offset = tl.minimum(offset, num_prefills - 1)
-
-    seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask)
-    qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask)
-    qsl_end = tl.load(query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask)
-
-    query_len = qsl_end - qsl_start
-    prefix_len = seq_len - query_len
-    gather_len = query_len + tl.minimum(prefix_len, window_size - 1)
-
-    tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask)
-
-
@triton.jit(do_not_specialize=["token_offset"])
 def _compute_swa_indices_and_lens_kernel(
    swa_indices_ptr,
@@ -939,19 +939,23 @@ def get_max_concurrency_for_kv_cache_config(
 ) -> float:
    """
    Get the maximum concurrency for the given KV cache configuration.
+
+    A request at max_model_len consumes whole blocks from each group's block
+    table — cdiv(per-request bytes, page bytes) of the group's spec — and all
+    groups draw those block ids from one shared pool, so the per-request
+    total is the sum over groups. The memory/page ratio is identical whether
+    a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or
+    a representative per-layer spec (scheduler config), so both capacity
+    call sites agree.
    """
-    num_layer_per_group = max(
-        len(group.layer_names) for group in kv_cache_config.kv_cache_groups
+    num_blocks_per_request = sum(
+        cdiv(
+            group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
+            group.kv_cache_spec.page_size_bytes,
+        )
+        for group in kv_cache_config.kv_cache_groups
    )
-    max_memory_usage_per_request = num_layer_per_group * max_memory_usage_bytes(
-        vllm_config, (group.kv_cache_spec for group in kv_cache_config.kv_cache_groups)
-    )
-    memory_per_block = (
-        kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
-        * num_layer_per_group
-    )
-    num_block_per_request = cdiv(max_memory_usage_per_request, memory_per_block)
-    max_concurrency = kv_cache_config.num_blocks / num_block_per_request
+    max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request
    return max_concurrency


--- 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
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