Sparse MLA on Hopper: Use SGLang's kernel for the sparse mla low latency runs

Signed-off-by: Alexander Matveev <amatveev@redhat.com>

.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh

@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
   exit 1
 fi
 
-#echo "--- DP+TP"
-#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
-#server_pid=$!
-#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
-#vllm bench serve \
-#    --backend vllm \
-#    --dataset-name random \
-#    --model meta-llama/Llama-3.2-3B-Instruct \
-#    --num-prompts 20 \
-#    --result-dir ./test_results \
-#    --result-filename dp_pp.json \
-#    --save-result \
-#    --endpoint /v1/completions
-#kill -s SIGTERM $server_pid; wait $server_pid || true
-#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
-#if [ "$failed_req" -ne 0 ]; then
-#  echo "Some requests were failed!"
-#  exit 1
-#fi
+echo "--- DP+TP"
+vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
+server_pid=$!
+timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
+vllm bench serve \
+   --backend vllm \
+   --dataset-name random \
+   --model meta-llama/Llama-3.2-3B-Instruct \
+   --num-prompts 20 \
+   --result-dir ./test_results \
+   --result-filename dp_pp.json \
+   --save-result \
+   --endpoint /v1/completions
+kill -s SIGTERM $server_pid; wait $server_pid || true
+failed_req=$(jq '.failed' ./test_results/dp_pp.json)
+if [ "$failed_req" -ne 0 ]; then
+ echo "Some requests were failed!"
+ exit 1
+fi

.github/CODEOWNERS

@@ -44,9 +44,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /vllm/pooling_params.py @noooop @DarkLight1337
 /vllm/tokenizers @DarkLight1337 @njhill
 /vllm/renderers @DarkLight1337 @njhill
-/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33
-/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33
-/vllm/parser @aarnphm @chaunceyjiang @sfeng33
+/vllm/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/vllm/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/vllm/parser @aarnphm @chaunceyjiang @sfeng33 @bbrowning
 
 # vLLM V1
 /vllm/v1/attention @LucasWilkinson @MatthewBonanni
@@ -93,9 +93,9 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
 /tests/v1/kv_connector @ApostaC @orozery
 /tests/v1/kv_offload @ApostaC @orozery
 /tests/v1/determinism @yewentao256
-/tests/reasoning @aarnphm @chaunceyjiang @sfeng33
-/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33
-/tests/tool_use @aarnphm @chaunceyjiang @sfeng33
+/tests/reasoning @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/tests/tool_parsers @aarnphm @chaunceyjiang @sfeng33 @bbrowning
+/tests/tool_use @aarnphm @chaunceyjiang @sfeng33 @bbrowning
 
 # Transformers modeling backend
 /vllm/model_executor/models/transformers @hmellor

.github/workflows/macos-smoke-test.yml

@@ -45,6 +45,7 @@ jobs:
       - name: Smoke test vllm serve
         run: |
           # Start server in background
+          VLLM_CPU_KVCACHE_SPACE=1 \
           vllm serve Qwen/Qwen3-0.6B \
             --max-model-len=2K \
             --load-format=dummy \

.github/workflows/new_pr_bot.yml

@@ -62,14 +62,14 @@ jobs:
             const prAuthor = context.payload.pull_request.user.login;
 
             const { data: searchResults } = await github.rest.search.issuesAndPullRequests({
-              q: `repo:${owner}/${repo} type:pr author:${prAuthor}`,
+              q: `repo:${owner}/${repo} type:pr is:merged author:${prAuthor}`,
               per_page: 1,
             });
 
-            const authorPRCount = searchResults.total_count;
-            console.log(`Found ${authorPRCount} PRs by ${prAuthor}`);
+            const mergedPRCount = searchResults.total_count;
+            console.log(`Found ${mergedPRCount} merged PRs by ${prAuthor}`);
 
-            if (authorPRCount === 1) {
+            if (mergedPRCount === 0) {
               console.log(`Posting welcome comment for first-time contributor: ${prAuthor}`);
               await github.rest.issues.createComment({
                 owner,
@@ -98,5 +98,5 @@ jobs:
                 ].join('\n'),
               });
             } else {
-              console.log(`Skipping comment for ${prAuthor} - not their first PR (${authorPRCount} PRs found)`);
+              console.log(`Skipping comment for ${prAuthor} - not a first-time contributor (${mergedPRCount} merged PRs)`);
             }

CMakeLists.txt

@@ -923,6 +923,14 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    # nvfp4_kv_cache_kernels uses non-stable torch API and is called directly
+    # from cache_kernels.cu, so it belongs in _C rather than _C_stable.
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${NVFP4_KV_SRC}"
+      CUDA_ARCHS "${FP4_ARCHS}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM120=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM120=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM120=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -949,6 +957,12 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       SRCS "${SRCS}"
       CUDA_ARCHS "${FP4_ARCHS}")
     list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
+    set(NVFP4_KV_SRC "csrc/nvfp4_kv_cache_kernels.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${NVFP4_KV_SRC}"
+      CUDA_ARCHS "${FP4_ARCHS}")
+    target_sources(_C PRIVATE ${NVFP4_KV_SRC})
+    target_compile_definitions(_C PRIVATE ENABLE_NVFP4_SM100=1)
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_NVFP4_SM100=1")
     list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM100=1")
     message(STATUS "Building NVFP4 for archs: ${FP4_ARCHS}")
@@ -1226,6 +1240,7 @@ endif()
 if (VLLM_GPU_LANG STREQUAL "CUDA")
     include(cmake/external_projects/deepgemm.cmake)
     include(cmake/external_projects/flashmla.cmake)
+    include(cmake/external_projects/cutlass_fa3.cmake)
     include(cmake/external_projects/qutlass.cmake)
 
     # vllm-flash-attn should be last as it overwrites some CMake functions

cmake/cpu_extension.cmake

@@ -30,6 +30,21 @@ else()
     list(APPEND CXX_COMPILE_FLAGS
         "-fopenmp"
         "-DVLLM_CPU_EXTENSION")
+
+    # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
+    # and create a local shim dir with it
+    vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
+
+    find_library(OPEN_MP
+        NAMES gomp
+        PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
+        NO_DEFAULT_PATH
+        REQUIRED
+    )
+    # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
+    if (OPEN_MP)
+        set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
+    endif()
 endif()
 
 if (NOT MACOSX_FOUND)
@@ -175,20 +190,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
         if(NOT NPROC)
             set(NPROC 4)
         endif()
-        # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
-        # and create a local shim dir with it
-        vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
-
-        find_library(OPEN_MP
-            NAMES gomp
-            PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
-            NO_DEFAULT_PATH
-            REQUIRED
-        )
-        # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
-        if (OPEN_MP)
-            set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
-        endif()
 
         # Fetch and populate ACL
         if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")

cmake/external_projects/cutlass_fa3.cmake

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# CUTLASS FA3 MLA Sparse Attention — requires CUDA >= 12.4, SM90a
+#
+# Vendors the sgl-attn CUTLASS FlashAttention3 kernel from SGLang into vLLM
+# as a self-contained extension (_cutlass_fa3_C). This provides a high-
+# performance sparse MLA attention kernel for SM90 (Hopper) GPUs.
+#
+# Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+# CUTLASS: https://github.com/NVIDIA/cutlass (commit 57e3cfb4)
+
+# Guard: CUDA >= 12.4 required for SM90a features used by FA3
+if(NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL "12.4")
+  message(STATUS "Skipping CUTLASS FA3: requires CUDA >= 12.4")
+  # Create empty target so setup.py doesn't fail on unsupported systems
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+# Guard: SM90 architecture required
+set(CUTLASS_FA3_SUPPORT_ARCHS)
+list(APPEND CUTLASS_FA3_SUPPORT_ARCHS "9.0a")
+cuda_archs_loose_intersection(
+  CUTLASS_FA3_ARCHS "${CUTLASS_FA3_SUPPORT_ARCHS}" "${CUDA_ARCHS}")
+if(NOT CUTLASS_FA3_ARCHS)
+  message(STATUS "Skipping CUTLASS FA3: requires SM90 (CUDA_ARCHS=${CUDA_ARCHS})")
+  add_custom_target(_cutlass_fa3_C)
+  return()
+endif()
+
+include(FetchContent)
+
+# Fetch sgl-attn (Flash Attention 3 kernels from SGLang)
+# We only need the source files, not the build system, so we use
+# FetchContent_Populate to download without building.
+if (DEFINED ENV{SGL_ATTN_SRC_DIR})
+  set(SGL_ATTN_SRC_DIR $ENV{SGL_ATTN_SRC_DIR})
+endif()
+if(SGL_ATTN_SRC_DIR)
+  FetchContent_Declare(cutlass_fa3
+    SOURCE_DIR ${SGL_ATTN_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_fa3
+    GIT_REPOSITORY https://github.com/sgl-project/sgl-attn.git
+    GIT_TAG bcf72ccc6816b36a5fae2c5a3c027604629785e0
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_fa3)
+if(NOT cutlass_fa3_POPULATED)
+  FetchContent_Populate(cutlass_fa3)
+endif()
+message(STATUS "CUTLASS FA3 sgl-attn source: ${cutlass_fa3_SOURCE_DIR}")
+
+# Fetch CUTLASS for FA3 (headers only, separate from vLLM's main CUTLASS
+# to avoid version conflicts). Use FetchContent_Populate to avoid running
+# CUTLASS's own CMakeLists.txt which would create conflicting targets.
+if (DEFINED ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+  set(CUTLASS_FA3_CUTLASS_SRC_DIR $ENV{CUTLASS_FA3_CUTLASS_SRC_DIR})
+endif()
+if(CUTLASS_FA3_CUTLASS_SRC_DIR)
+  FetchContent_Declare(cutlass_for_fa3
+    SOURCE_DIR ${CUTLASS_FA3_CUTLASS_SRC_DIR})
+else()
+  FetchContent_Declare(cutlass_for_fa3
+    GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
+    GIT_TAG 57e3cfb47a2d9e0d46eb6335c3dc411498efa198
+    GIT_PROGRESS TRUE
+    GIT_SHALLOW FALSE)
+endif()
+FetchContent_GetProperties(cutlass_for_fa3)
+if(NOT cutlass_for_fa3_POPULATED)
+  FetchContent_Populate(cutlass_for_fa3)
+endif()
+message(STATUS "CUTLASS FA3 cutlass source: ${cutlass_for_fa3_SOURCE_DIR}")
+
+set(FA3_SRC "${cutlass_fa3_SOURCE_DIR}/hopper")
+
+# flash_api.cpp dispatches to all head dimensions + dtypes (BF16, FP16, FP8)
+# at compile time. With FLASHATTENTION_DISABLE_SM8x, only SM90 instantiations
+# are needed. We exclude hdimall_* (fails on CUDA 13+) and backward files.
+file(GLOB FA3_INSTANTIATION_SOURCES
+  # BF16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_bf16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_bf16*_sm90.cu"
+  # FP16 instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_fp16*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_fp16*_sm90.cu"
+  # FP8 (e4m3) instantiations
+  "${FA3_SRC}/instantiations/flash_fwd_hdim64_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim96_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim128_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim192_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdim256_e4m3*_sm90.cu"
+  "${FA3_SRC}/instantiations/flash_fwd_hdimdiff_e4m3*_sm90.cu")
+
+set(FA3_CORE_SOURCES
+  "${FA3_SRC}/flash_api.cpp"
+  "${FA3_SRC}/flash_prepare_scheduler.cu"
+  "${FA3_SRC}/flash_fwd_combine.cu")
+
+set(FA3_ALL_SOURCES
+  "${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_fa3_extension.cc"
+  ${FA3_CORE_SOURCES}
+  ${FA3_INSTANTIATION_SOURCES})
+
+set(FA3_INCLUDE_DIRS
+  ${FA3_SRC}
+  ${cutlass_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/include
+  ${cutlass_for_fa3_SOURCE_DIR}/tools/util/include
+  ${CMAKE_CURRENT_SOURCE_DIR}/csrc)
+
+# Set SM90a gencode flags for all FA3 CUDA sources
+set_gencode_flags_for_srcs(
+  SRCS "${FA3_ALL_SOURCES}"
+  CUDA_ARCHS "${CUTLASS_FA3_ARCHS}")
+
+define_extension_target(_cutlass_fa3_C
+  DESTINATION vllm
+  LANGUAGE ${VLLM_GPU_LANG}
+  SOURCES ${FA3_ALL_SOURCES}
+  COMPILE_FLAGS ${VLLM_GPU_FLAGS}
+  ARCHITECTURES ${VLLM_GPU_ARCHES}
+  INCLUDE_DIRECTORIES ${FA3_INCLUDE_DIRS}
+  USE_SABI 3
+  WITH_SOABI)
+
+# FA3-specific compile options for CUDA and C++ source files:
+# - C++17 required by CUTLASS
+# - Fast math for performance
+# - Relaxed constexpr for CUTLASS template metaprogramming
+# - Disable backward pass, dropout, uneven K (not needed for inference)
+# - Enable varlen-only mode (all our use cases are variable-length)
+target_compile_options(_cutlass_fa3_C PRIVATE
+  $<$<COMPILE_LANGUAGE:CUDA>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CXX>:-UPy_LIMITED_API>
+  $<$<COMPILE_LANGUAGE:CUDA>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CXX>:-std=c++17>
+  $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>
+  $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
+
+target_compile_definitions(_cutlass_fa3_C PRIVATE
+  CUTE_USE_PACKED_TUPLE=1
+  CUTLASS_ENABLE_GDC_FOR_SM90
+  CUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED
+  CUTLASS_ENABLE_TENSOR_CORE_MMA=1
+  FLASHATTENTION_DISABLE_BACKWARD
+  FLASHATTENTION_DISABLE_DROPOUT
+  FLASHATTENTION_DISABLE_UNEVEN_K
+  FLASHATTENTION_DISABLE_SM8x
+  FLASHATTENTION_VARLEN_ONLY)
+
+message(STATUS "CUTLASS FA3 MLA Sparse: enabled for SM90 (${CUTLASS_FA3_ARCHS})")

csrc/cache_kernels.cu

@@ -724,6 +724,28 @@ void reshape_and_cache_flash(
   int num_tokens = slot_mapping.size(0);
   int num_heads = key.size(1);
   int head_size = key.size(2);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  if (kv_cache_dtype == "nvfp4") {
+#if defined(ENABLE_NVFP4_SM100) || defined(ENABLE_NVFP4_SM120)
+    // NVFP4 dispatch is compiled separately for SM100+.
+    extern void reshape_and_cache_nvfp4_dispatch(
+        torch::Tensor & key, torch::Tensor & value, torch::Tensor & key_cache,
+        torch::Tensor & value_cache, torch::Tensor & slot_mapping,
+        torch::Tensor & k_scale, torch::Tensor & v_scale);
+    reshape_and_cache_nvfp4_dispatch(key, value, key_cache, value_cache,
+                                     slot_mapping, k_scale, v_scale);
+    return;
+#else
+    TORCH_CHECK(false,
+                "NVFP4 KV cache requires SM100+ (Blackwell). "
+                "Please rebuild vllm with a Blackwell-compatible CUDA target.");
+#endif
+  }
+
+  // Original FP8/auto path.
   int block_size = key_cache.size(1);
 
   int64_t key_stride = key.stride(0);
@@ -741,8 +763,6 @@ void reshape_and_cache_flash(
 
   dim3 grid(num_tokens);
   dim3 block(std::min(num_heads * head_size, 512));
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
-  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
   DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
                              CALL_RESHAPE_AND_CACHE_FLASH);

csrc/cpu/torch_bindings.cpp

@@ -141,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
                                       torch::Tensor slot_mapping,
                                       const int64_t block_size);
 
+void init_cpu_memory_env(std::vector<int64_t> node_ids);
+
 namespace cpu_utils {
 void eagle_prepare_inputs_padded_kernel_impl(
     const torch::Tensor& cu_num_draft_tokens,
@@ -431,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "block_size) -> ()",
       &compute_slot_mapping_kernel_impl);
 
+  ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
+
   // Speculative decoding kernels
   ops.def(
       "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "

csrc/cpu/utils.cpp

@@ -13,13 +13,80 @@
 #include "cpu/utils.hpp"
 
 #ifdef VLLM_NUMA_DISABLED
-std::string init_cpu_threads_env(const std::string& cpu_ids) {
-  return std::string(
-      "Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
-      "no effect to setup thread affinity.");
-}
+void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
+#else
+void init_cpu_memory_env(std::vector<int64_t> node_ids) {
+  // Memory node binding
+  if (numa_available() != -1) {
+    // Concatenate all node_ids into a single comma-separated string
+    if (!node_ids.empty()) {
+      std::string node_ids_str;
+      for (const int node_id : node_ids) {
+        if (!node_ids_str.empty()) {
+          node_ids_str += ",";
+        }
+        node_ids_str += std::to_string(node_id);
+      }
 
-#endif
+      bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
+      bitmask* src_mask = numa_get_mems_allowed();
+
+      int pid = getpid();
+
+      if (mask && src_mask) {
+        // move all existing pages to the specified numa node.
+        *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
+        int page_num = numa_migrate_pages(pid, src_mask, mask);
+        if (page_num == -1) {
+          TORCH_WARN("numa_migrate_pages failed. errno: " +
+                     std::to_string(errno));
+        }
+
+        // Restrict memory allocation to the selected NUMA node(s).
+        // Enhances memory locality for the threads bound to those NUMA CPUs.
+        if (node_ids.size() > 1) {
+          errno = 0;
+          numa_set_interleave_mask(mask);
+          if (errno != 0) {
+            TORCH_WARN("numa_set_interleave_mask failed. errno: " +
+                       std::to_string(errno));
+          } else {
+            TORCH_WARN(
+                "NUMA binding: Using INTERLEAVE policy for memory "
+                "allocation across multiple NUMA nodes (nodes: " +
+                node_ids_str +
+                "). Memory allocations will be "
+                "interleaved across the specified NUMA nodes.");
+          }
+        } else {
+          errno = 0;
+          numa_set_membind(mask);
+          if (errno != 0) {
+            TORCH_WARN("numa_set_membind failed. errno: " +
+                       std::to_string(errno));
+          } else {
+            TORCH_WARN(
+                "NUMA binding: Using MEMBIND policy for memory "
+                "allocation on the NUMA nodes (" +
+                node_ids_str +
+                "). Memory allocations will be "
+                "strictly bound to these NUMA nodes.");
+          }
+        }
+
+        numa_set_strict(1);
+
+        numa_free_nodemask(mask);
+        numa_free_nodemask(src_mask);
+      } else {
+        TORCH_WARN(
+            "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
+            std::to_string(errno));
+      }
+    }
+  }
+}
+#endif  // VLLM_NUMA_DISABLED
 
 namespace cpu_utils {
 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {

csrc/cutlass_fa3_extension.cc

@@ -0,0 +1,72 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * Vendored CUTLASS FA3 MLA attention kernel binding for vLLM.
+ * Based on sgl-kernel/csrc/flash_extension.cc from SGLang.
+ *
+ * This registers the FA3 forward pass as a PyTorch C++ extension under
+ * the _cutlass_fa3_C namespace, enabling torch.ops._cutlass_fa3_C.fwd().
+ *
+ * Original source:
+ *   https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+ *   sgl-kernel/csrc/flash_extension.cc
+ */
+#include <Python.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/all.h>
+#include <torch/library.h>
+
+#include "sgl_flash_kernel_ops.h"
+
+TORCH_LIBRARY_FRAGMENT(_cutlass_fa3_C, m) {
+  /*
+   * CUTLASS FA3 MLA forward pass.
+   * Signature matches sgl-attn's mha_fwd() exactly.
+   */
+  m.def(
+      "fwd(Tensor   q,"
+      "    Tensor   k,"
+      "    Tensor   v,"
+      "    Tensor?  k_new,"
+      "    Tensor?  v_new,"
+      "    Tensor?  q_v,"
+      "    Tensor?  out,"
+      "    Tensor?  cu_seqlens_q,"
+      "    Tensor?  cu_seqlens_k,"
+      "    Tensor?  cu_seqlens_k_new,"
+      "    Tensor?  seqused_q,"
+      "    Tensor?  seqused_k,"
+      "    int?     max_seqlen_q,"
+      "    int?     max_seqlen_k,"
+      "    Tensor?  page_table,"
+      "    Tensor?  kv_batch_idx,"
+      "    Tensor?  leftpad_k,"
+      "    Tensor?  rotary_cos,"
+      "    Tensor?  rotary_sin,"
+      "    Tensor?  seqlens_rotary,"
+      "    Tensor?  q_descale,"
+      "    Tensor?  k_descale,"
+      "    Tensor?  v_descale,"
+      "    float?   softmax_scale,"
+      "    bool     is_causal,"
+      "    int      window_size_left,"
+      "    int      window_size_right,"
+      "    int      attention_chunk,"
+      "    float    softcap,"
+      "    bool     is_rotary_interleaved,"
+      "    Tensor?  scheduler_metadata,"
+      "    int      num_splits,"
+      "    bool?    pack_gqa,"
+      "    int      sm_margin,"
+      "    Tensor?  sinks"
+      ") -> (Tensor, Tensor, Tensor, Tensor)");
+
+  m.impl("fwd", torch::kCUDA, make_pytorch_shim(&mha_fwd));
+}
+
+// Python module initialization for _cutlass_fa3_C
+PyMODINIT_FUNC PyInit__cutlass_fa3_C() {
+  static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, "_cutlass_fa3_C",
+                                      nullptr, 0, nullptr};
+  return PyModule_Create(&module);
+}

csrc/nvfp4_kv_cache_kernels.cu

@@ -0,0 +1,275 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+// NVFP4 KV cache store kernel.
+// Quantizes bf16 key/value to packed FP4 + FP8 block scales and writes them
+// into the paged KV cache.
+//
+// Per page layout: [K_data | K_scale | V_data | V_scale]
+// Both data and scale regions are contiguous per head, enabling direct
+// TMA descriptor use.
+//
+// Reuses device functions from nvfp4_utils.cuh:
+//   - cvt_warp_fp16_to_fp4()  for bf16 → fp4 quantization + block scale
+//   - pack_fp4()              for packing float pairs to fp4
+//   - reciprocal_approximate_ftz() for fast reciprocal
+
+#define NVFP4_ENABLE_ELTS16 1
+#include "libtorch_stable/quantization/fp4/nvfp4_utils.cuh"
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+#include "dispatch_utils.h"
+
+namespace vllm {
+
+// Compute swizzled scale offset for SM100 trtllm-gen MHA kernel.
+// The swizzle pattern for HND layout is:
+//   [T//4, 4, 4, S//4] → permute(0, 2, 3, 1) → reshape to [T, S]
+// where T = block_size (page_size), S = scale_dim = head_size // 16.
+//
+// For a linear (t, s) position, the swizzled position is:
+//   swizzled_t = (t / 4) * 4 + (s / (S / 4))
+//   swizzled_s = (s % (S / 4)) * 4 + (t % 4)
+__device__ __forceinline__ int swizzle_scale_offset(int t, int s,
+                                                    int scale_dim) {
+  int s_group = scale_dim / 4;
+  int swizzled_t = (t / 4) * 4 + (s / s_group);
+  int swizzled_s = (s % s_group) * 4 + (t % 4);
+  return swizzled_t * scale_dim + swizzled_s;
+}
+
+// Kernel: quantize bf16 key/value to NVFP4 and store in paged KV cache.
+//
+// Takes separate data and scale cache pointers for K and V.
+// Within each KV side, data and scale are separate contiguous regions.
+//
+// Threading: one CUDA block per token, threads process heads and
+// groups of 16 elements within each head.
+template <typename scalar_t>
+__global__ void reshape_and_cache_nvfp4_kernel(
+    const scalar_t* __restrict__ key,      // [num_tokens, num_heads, head_size]
+    const scalar_t* __restrict__ value,    // [num_tokens, num_heads, head_size]
+    uint8_t* __restrict__ key_data_cache,  // data region for K
+    uint8_t* __restrict__ value_data_cache,    // data region for V
+    uint8_t* __restrict__ key_scale_cache,     // scale region for K
+    uint8_t* __restrict__ value_scale_cache,   // scale region for V
+    const int64_t* __restrict__ slot_mapping,  // [num_actual_tokens]
+    const float* __restrict__ k_scale_ptr,     // pointer to checkpoint k_scale
+    const float* __restrict__ v_scale_ptr,     // pointer to checkpoint v_scale
+    const int64_t key_stride,                  // key.stride(0) in elements
+    const int64_t value_stride,                // value.stride(0) in elements
+    const int num_heads, const int head_size, const int block_size,
+    const int64_t data_block_stride,         // data cache stride for dim 0
+    const int64_t data_head_stride,          // data cache stride for heads
+    const int64_t data_block_offset_stride,  // data cache stride for tokens
+    const int64_t scale_block_stride,        // scale cache stride for dim 0
+    const int64_t scale_head_stride,         // scale cache stride for heads
+    const int64_t scale_block_offset_stride  // scale cache stride for tokens
+) {
+  using CudaType = typename CUDATypeConverter<scalar_t>::Type;
+  using PVec = PackedVec<CudaType, CVT_FP4_PACK16>;
+
+  static constexpr int ELTS = CVT_FP4_ELTS_PER_THREAD;  // 16 or 8
+  static constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / ELTS;
+
+  const int64_t token_idx = blockIdx.x;
+  const int64_t slot_idx = slot_mapping[token_idx];
+  if (slot_idx < 0) return;
+
+  const int64_t block_idx = slot_idx / block_size;
+  const int block_offset = static_cast<int>(slot_idx % block_size);
+
+  const int scale_dim = head_size / 16;
+  const int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
+
+  const int total_groups = num_heads * groups_per_head;
+  const int tid = threadIdx.x;
+  const int num_thread_groups = blockDim.x / THREADS_PER_SF;
+  const int tg_id = tid / THREADS_PER_SF;
+  const int tg_lane = tid % THREADS_PER_SF;
+
+  // Process both K (kv=0) and V (kv=1)
+#pragma unroll
+  for (int kv = 0; kv < 2; kv++) {
+    const scalar_t* __restrict__ src = (kv == 0) ? key : value;
+    const float global_scale = 1.0f / ((kv == 0) ? *k_scale_ptr : *v_scale_ptr);
+    const int64_t src_stride = (kv == 0) ? key_stride : value_stride;
+    uint8_t* __restrict__ data_cache =
+        (kv == 0) ? key_data_cache : value_data_cache;
+    uint8_t* __restrict__ sc_cache =
+        (kv == 0) ? key_scale_cache : value_scale_cache;
+
+    // Source pointer for this token (use actual stride, not assumed contiguous)
+    const CudaType* __restrict__ token_src =
+        reinterpret_cast<const CudaType*>(src) + token_idx * src_stride;
+
+    // Destination bases in data and scale caches for this token's block
+    uint8_t* __restrict__ data_block =
+        data_cache + block_idx * data_block_stride;
+    uint8_t* __restrict__ scale_block =
+        sc_cache + block_idx * scale_block_stride;
+
+    for (int g = tg_id; g < total_groups; g += num_thread_groups) {
+      const int head = g / groups_per_head;
+      const int group_in_head = g % groups_per_head;
+
+      // Load 16 (or 8) bf16 elements from source
+      PVec in_vec;
+      const CudaType* __restrict__ src_ptr =
+          token_src + head * head_size + group_in_head * CVT_FP4_SF_VEC_SIZE +
+          tg_lane * ELTS;
+
+#pragma unroll
+      for (int i = 0; i < ELTS / 2; i++) {
+        in_vec.elts[i] = reinterpret_cast<
+            const typename PackedTypeConverter<CudaType>::Type*>(src_ptr)[i];
+      }
+
+      // Quantize: produces packed fp4 and writes scale factor.
+      uint8_t sf_val;
+      uint8_t* sf_out_ptr = (tg_lane == 0) ? &sf_val : nullptr;
+
+      fp4_packed_t packed = cvt_warp_fp16_to_fp4<CudaType, THREADS_PER_SF>(
+          in_vec, global_scale, sf_out_ptr);
+
+      // Write packed FP4 data to data cache
+      uint8_t* __restrict__ data_dst = data_block + head * data_head_stride +
+                                       block_offset * data_block_offset_stride;
+
+#if CVT_FP4_PACK16
+      {
+        // 16 elements → 8 bytes (u32x2)
+        int data_byte_offset = group_in_head * 8;
+        reinterpret_cast<uint64_t*>(data_dst + data_byte_offset)[0] =
+            (uint64_t(packed.hi) << 32) | uint64_t(packed.lo);
+      }
+#else
+      {
+        // 8 elements → 4 bytes (uint32_t)
+        int data_byte_offset =
+            group_in_head * CVT_FP4_SF_VEC_SIZE / 2 + tg_lane * ELTS / 2;
+        reinterpret_cast<uint32_t*>(data_dst + data_byte_offset)[0] = packed;
+      }
+#endif
+
+      // Write block scale to scale cache.
+      // K (kv==0): linear layout (no swizzle).
+      // V (kv==1): swizzled layout for SM100 trtllm-gen MHA kernel.
+      if (sf_out_ptr != nullptr) {
+        int scale_idx = group_in_head;
+        uint8_t* __restrict__ scale_dst;
+        if (kv == 0) {
+          scale_dst = scale_block + head * scale_head_stride +
+                      block_offset * scale_block_offset_stride + scale_idx;
+        } else {
+          int swizzled_offset =
+              swizzle_scale_offset(block_offset, scale_idx, scale_dim);
+          int swizzled_t = swizzled_offset / scale_dim;
+          int swizzled_s = swizzled_offset % scale_dim;
+          scale_dst = scale_block + head * scale_head_stride +
+                      swizzled_t * scale_block_offset_stride + swizzled_s;
+        }
+        *scale_dst = sf_val;
+      }
+    }
+  }
+}
+
+}  // namespace vllm
+
+// Non-template entry point callable from cache_kernels.cu.
+// Receives key_cache/value_cache as kv_cache[:, 0] and kv_cache[:, 1].
+// Each KV side contains both data and scale:
+//   page = [K_data | K_scale | V_data | V_scale]
+void reshape_and_cache_nvfp4_dispatch(torch::Tensor& key, torch::Tensor& value,
+                                      torch::Tensor& key_cache,
+                                      torch::Tensor& value_cache,
+                                      torch::Tensor& slot_mapping,
+                                      torch::Tensor& k_scale,
+                                      torch::Tensor& v_scale) {
+  int num_tokens = slot_mapping.size(0);
+  int num_heads = key.size(1);
+  int head_size = key.size(2);
+  int data_dim = head_size / 2;
+  int scale_dim = head_size / 16;
+  int full_dim = data_dim + scale_dim;
+
+  // key_cache is kv_cache[:, 0] with shape
+  // [num_blocks, block_size, num_heads, full_dim] in logical order.
+  // Strides encode the physical layout (HND or NHD).
+  TORCH_CHECK(key_cache.dim() == 4, "key_cache must be 4D");
+  TORCH_CHECK(key_cache.size(3) == full_dim,
+              "key_cache last dim must be data_dim + scale_dim, got ",
+              key_cache.size(3), " expected ", full_dim);
+
+  int block_size = key_cache.size(1);
+
+  TORCH_CHECK(head_size % 16 == 0,
+              "head_size must be divisible by 16 for NVFP4 KV cache");
+  TORCH_CHECK(block_size % 4 == 0,
+              "block_size must be divisible by 4 for NVFP4 KV cache swizzle");
+
+  // Detect physical layout from strides (based on full_dim).
+  // HND: head stride > block_offset stride.
+  bool is_hnd = key_cache.stride(2) > key_cache.stride(1);
+
+  int64_t data_block_stride = key_cache.stride(0);  // page_bytes
+  int64_t data_head_stride, data_block_offset_stride;
+  if (is_hnd) {
+    data_head_stride = (int64_t)block_size * data_dim;
+    data_block_offset_stride = data_dim;
+  } else {
+    data_head_stride = data_dim;
+    data_block_offset_stride = (int64_t)num_heads * data_dim;
+  }
+
+  // Page layout: [K_data | K_scale | V_data | V_scale]
+  // Scale follows data within each KV side.
+  int64_t data_per_kv = (int64_t)num_heads * block_size * data_dim;
+
+  uint8_t* key_scale_ptr = key_cache.data_ptr<uint8_t>() + data_per_kv;
+  uint8_t* value_scale_ptr = value_cache.data_ptr<uint8_t>() + data_per_kv;
+
+  // Scale strides: same page stride, inner strides from layout.
+  int64_t scale_block_stride = data_block_stride;
+  int64_t scale_head_stride, scale_block_offset_stride;
+  if (is_hnd) {
+    scale_head_stride = (int64_t)block_size * scale_dim;
+    scale_block_offset_stride = scale_dim;
+  } else {
+    scale_head_stride = scale_dim;
+    scale_block_offset_stride = (int64_t)num_heads * scale_dim;
+  }
+
+  const float* k_scale_ptr = k_scale.data_ptr<float>();
+  const float* v_scale_ptr = v_scale.data_ptr<float>();
+
+  int groups_per_head = head_size / CVT_FP4_SF_VEC_SIZE;
+  int total_groups = num_heads * groups_per_head;
+  constexpr int THREADS_PER_SF = CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
+  int num_threads = std::min(total_groups * THREADS_PER_SF, 512);
+  num_threads = ((num_threads + 31) / 32) * 32;
+
+  dim3 grid(num_tokens);
+  dim3 block(num_threads);
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  AT_DISPATCH_REDUCED_FLOATING_TYPES(
+      key.scalar_type(), "reshape_and_cache_nvfp4", [&] {
+        vllm::reshape_and_cache_nvfp4_kernel<scalar_t>
+            <<<grid, block, 0, stream>>>(
+                key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
+                key_cache.data_ptr<uint8_t>(), value_cache.data_ptr<uint8_t>(),
+                key_scale_ptr, value_scale_ptr,
+                slot_mapping.data_ptr<int64_t>(), k_scale_ptr, v_scale_ptr,
+                key.stride(0), value.stride(0), num_heads, head_size,
+                block_size, data_block_stride, data_head_stride,
+                data_block_offset_stride, scale_block_stride, scale_head_stride,
+                scale_block_offset_stride);
+      });
+}

csrc/sgl_flash_kernel_ops.h

@@ -0,0 +1,45 @@
+/* SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_flash_kernel_ops.h (commit bcf72ccc).
+ * Declares the mha_fwd() C++ function signature for CUTLASS FA3 kernels.
+ * NO MODIFICATIONS from the original (except removing unused macros).
+ */
+
+#pragma once
+
+#include <ATen/ATen.h>
+#include <ATen/Tensor.h>
+#include <torch/library.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+#include "sgl_kernel_torch_shim.h"
+
+/*
+ * From flash-attention (sgl-attn fork)
+ */
+std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
+    at::Tensor q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
+    at::Tensor k,  // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
+    at::Tensor v,  // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
+    std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
+    std::optional<at::Tensor> q_v_,  // MLA value projection query
+    std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
+    std::optional<at::Tensor> cu_seqlens_k_,
+    std::optional<at::Tensor> cu_seqlens_k_new_,
+    std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
+    std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
+    std::optional<at::Tensor> page_table_,
+    std::optional<at::Tensor> kv_batch_idx_,
+    std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
+    std::optional<at::Tensor> rotary_sin_,
+    std::optional<at::Tensor> seqlens_rotary_,
+    std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
+    std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
+    bool is_causal, int64_t window_size_left, int64_t window_size_right,
+    int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
+    std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
+    std::optional<bool> pack_gqa_, int64_t sm_margin,
+    std::optional<const at::Tensor>& sinks_);

csrc/sgl_kernel_torch_shim.h

@@ -0,0 +1,121 @@
+/* Adapted from:
+ * https://github.com/neuralmagic/vllm-flash-attention/blob/90eacc1af2a7c3de62ea249e929ed5faccf38954/csrc/common/pytorch_shim.h
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright 2025 SGLang Team. All Rights Reserved.
+ *
+ * Vendored from sgl-kernel/include/sgl_kernel_torch_shim.h (commit bcf72ccc).
+ * Provides make_pytorch_shim() template for PyTorch op registration type
+ * conversion. NO MODIFICATIONS from the original.
+ */
+
+#pragma once
+
+#include <torch/library.h>
+
+/**
+ * Unfortunately, the type signatures of the flash_attn ops are not compatible
+ * with the PyTorch library bindings. To get around that we use
+ * `make_pytorch_shim` which creates a lambda that exposes the API using
+ * PyTorch compatible types to the types, then converts them to the types
+ * expected by the flash_attn ops. This shims allows us to make minimal changes
+ * to `flash_api.cpp` making it easier to synchronize with upstream changes.
+ *
+ * The `pytorch_library_compatible_type` struct is used to map from the
+ * flash_attn ops types to a PyTorch library compatible one. The main issues is
+ * that the following types are not support by PyTorch library bindings:
+ *  - `int`
+ *  - `float`
+ *  - `std::optional<T> &`
+ *  - `std::optional<const at::Tensor> &`
+ * So we convert them to (respectively):
+ *  - `int64_t`
+ *  - `double`
+ *  - `const std::optional<T>&`
+ *  - `const std::optional<at::Tensor>&`
+ */
+
+template <typename T>
+struct pytorch_library_compatible_type {
+  using type = T;
+  static T convert_from_type(T arg) { return arg; }
+};
+
+template <typename T>
+using pytorch_library_compatible_type_t =
+    typename pytorch_library_compatible_type<T>::type;
+
+template <typename T>
+T convert_from_pytorch_compatible_type(
+    pytorch_library_compatible_type_t<T> arg) {
+  return pytorch_library_compatible_type<T>::convert_from_type(arg);
+}
+
+// Map `c10::optional<T> &` -> `const c10::optional<T>&`
+//  (NOTE: this is bit unsafe but non of the ops in flash_attn mutate
+//   the optional container)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>&> {
+  using type = const c10::optional<T>&;
+  static c10::optional<T>& convert_from_type(const c10::optional<T>& arg) {
+    return const_cast<c10::optional<T>&>(arg);
+  }
+};
+
+// Map `c10::optional<T>` ->
+//          `c10::optional<pytorch_library_compatible_type_t<T>>`
+//  (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
+template <typename T>
+struct pytorch_library_compatible_type<c10::optional<T>> {
+  using type = c10::optional<pytorch_library_compatible_type_t<T>>;
+  static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(
+      c10::optional<T> arg) {
+    return arg;
+  }
+};
+
+// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
+template <>
+struct pytorch_library_compatible_type<c10::optional<const at::Tensor>&> {
+  using type = const c10::optional<at::Tensor>&;
+  static c10::optional<const at::Tensor>& convert_from_type(
+      const c10::optional<at::Tensor>& arg) {
+    return const_cast<c10::optional<const at::Tensor>&>(
+        reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
+  }
+};
+
+// Map `int` -> `int64_t`
+template <>
+struct pytorch_library_compatible_type<int> {
+  using type = int64_t;
+  static int convert_from_type(int64_t arg) {
+    TORCH_CHECK(arg <= std::numeric_limits<int>::max(),
+                "int64_t value is too large to be converted to int");
+    TORCH_CHECK(arg >= std::numeric_limits<int>::min(),
+                "int64_t value is too small to be converted to int");
+    return arg;
+  }
+};
+
+// Map `float` -> `double`
+template <>
+struct pytorch_library_compatible_type<float> {
+  using type = double;
+  static float convert_from_type(double arg) {
+    TORCH_CHECK(std::abs(arg) <= std::numeric_limits<float>::max(),
+                "double value is too large to be converted to float");
+    return arg;
+  }
+};
+
+//
+//  Shim Utils
+//
+
+template <typename Ret, typename... Args>
+auto make_pytorch_shim(Ret (*fun)(Args... args)) {
+  return [fun](pytorch_library_compatible_type_t<Args>... args) {
+    return fun(convert_from_pytorch_compatible_type<Args>(args)...);
+  };
+}

docker/Dockerfile.cpu

@@ -173,7 +173,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
 
 RUN --mount=type=cache,target=/root/.cache/uv \
-    uv pip install -r requirements/dev.txt && \
+    uv pip install -r requirements/lint.txt && \
+    uv pip install -r requirements/test/cpu.txt && \
     pre-commit install --hook-type pre-commit --hook-type commit-msg
 
 ENTRYPOINT ["bash"]

docs/design/attention_backends.md

@@ -106,6 +106,7 @@ Priority is **1 = highest** (tried first).
 | 2 | `FLASH_ATTN` |
 | 3 | `TRITON_ATTN` |
 | 4 | `FLEX_ATTENTION` |
+| 5 | `TURBOQUANT` |
 
 **Ampere/Hopper (SM 8.x-9.x):**
 
@@ -115,6 +116,7 @@ Priority is **1 = highest** (tried first).
 | 2 | `FLASHINFER` |
 | 3 | `TRITON_ATTN` |
 | 4 | `FLEX_ATTENTION` |
+| 5 | `TURBOQUANT` |
 
 ### MLA Attention (DeepSeek-style)
 
@@ -130,16 +132,6 @@ Priority is **1 = highest** (tried first).
 | 6 | `FLASHINFER_MLA_SPARSE`**\*** |
 | 7 | `FLASHMLA_SPARSE` |
 
-**Ampere/Hopper (SM 8.x-9.x):**
-
-| Priority | Backend |
-| -------- | ------- |
-| 1 | `FLASH_ATTN_MLA` |
-| 2 | `FLASHMLA` |
-| 3 | `FLASHINFER_MLA` |
-| 4 | `TRITON_MLA` |
-| 5 | `FLASHMLA_SPARSE` |
-
 > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -207,6 +199,7 @@ configuration.
 
 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
+| `CUTLASS_FA3_MLA_SPARSE` | bf16 | `auto` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |

docs/design/metrics.md

@@ -42,7 +42,7 @@ These are documented under [Inferencing and Serving -> Production Metrics](../us
 
 ### Grafana Dashboard
 
-vLLM also provides [a reference example](../../examples/online_serving/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
+vLLM also provides [a reference example](../../examples/observability/prometheus_grafana/README.md) for how to collect and store these metrics using Prometheus and visualize them using a Grafana dashboard.
 
 The subset of metrics exposed in the Grafana dashboard gives us an indication of which metrics are especially important:
 
@@ -657,7 +657,7 @@ vLLM has support for OpenTelemetry tracing:
 - Added by <https://github.com/vllm-project/vllm/pull/4687> and reinstated by <https://github.com/vllm-project/vllm/pull/20372>
 - Configured with `--oltp-traces-endpoint` and `--collect-detailed-traces`
 - [OpenTelemetry blog post](https://opentelemetry.io/blog/2024/llm-observability/)
-- [User-facing docs](../../examples/online_serving/opentelemetry/README.md)
+- [User-facing docs](../../examples/observability/opentelemetry/README.md)
 - [Blog post](https://medium.com/@ronen.schaffer/follow-the-trail-supercharging-vllm-with-opentelemetry-distributed-tracing-aa655229b46f)
 - [IBM product docs](https://www.ibm.com/docs/en/instana-observability/current?topic=mgaa-monitoring-large-language-models-llms-vllm-public-preview)
 

docs/governance/committers.md

@@ -14,6 +14,7 @@ Sorted alphabetically by GitHub handle:
 - [@aarnphm](https://github.com/aarnphm): Structured output
 - [@alexm-redhat](https://github.com/alexm-redhat): Performance
 - [@ApostaC](https://github.com/ApostaC): Connectors, offloading
+- [@bbrowning](https://github.com/bbrowning): Tool use and reasoning parser
 - [@benchislett](https://github.com/benchislett): Engine core and spec decode
 - [@bigPYJ1151](https://github.com/bigPYJ1151): Intel CPU/XPU integration
 - [@chaunceyjiang](https://github.com/chaunceyjiang): Tool use and reasoning parser
@@ -121,7 +122,7 @@ If you have PRs touching the area, please feel free to ping the area owner for r
 - State space models: The state space models implementation in vLLM
     - @tdoublep, @tlrmchlsmth
 - Reasoning and tool calling parsers
-    - @chaunceyjiang, @aarnphm, @sfeng33
+    - @chaunceyjiang, @aarnphm, @sfeng33, @bbrowning
 
 ### Entrypoints
 

docs/models/supported_models.md

@@ -400,6 +400,7 @@ th {
 | `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ |
 | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ |
 | `Gemma3nForCausalLM` | Gemma 3n | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
+| `Gemma4ForCausalLM` | Gemma 4 | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |
 | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ |
 | `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ |
 | `Glm4MoeForCausalLM` | GLM-4.5, GLM-4.6, GLM-4.7 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ |
@@ -554,6 +555,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
 | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ |
 | `Gemma3ForConditionalGeneration` | Gemma 3 | T + I<sup>E+</sup> | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |
 | `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | |
+| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I<sup>+</sup> + V + A<sup>*</sup> | `google/gemma-4-E2B-it`, etc. | | ✅︎ |
 | `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
 | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
 | `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
@@ -633,6 +635,7 @@ Some models are supported only via the [Transformers modeling backend](#transfor
 <sup>^</sup> You need to set the architecture name via `--hf-overrides` to match the one in vLLM.</br>
 <sup>E</sup> Pre-computed embeddings can be inputted for this modality.</br>
 <sup>+</sup> Multiple items can be inputted per text prompt for this modality.
+<sup>*</sup> Only specific variants of the model support this modality (see notes below).</br>
 
 !!! note
     `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
@@ -643,6 +646,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor
     - Both audio and vision MM encoders use `transformers.AutoModel` implementation.  
     - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
 
+!!! note
+    For `Gemma4ForConditionalGeneration`:
+    - audio input is only supported by the `gemma-4-E2B` and `gemma-4-E4B` variants.
+    - The model does not ingest videos directly. However, vLLM’s Gemma 4 implementation supports video inputs by handling video processing internally. Users can send videos directly in the message structure to vLLM, where they are converted into text and image frames before being passed to the model.
+
 !!! note
     For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
 

examples/basic/offline_inference/basic.py

@@ -16,7 +16,7 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
 def main():
     # Create an LLM.
-    llm = LLM(model="nvidia/DeepSeek-V3.2-NVFP4", enforce_eager=True, tensor_parallel_size=4, kernel_config={"enable_flashinfer_autotune": False})
+    llm = LLM(model="facebook/opt-125m")
     # Generate texts from the prompts.
     # The output is a list of RequestOutput objects
     # that contain the prompt, generated text, and other information.

examples/observability/dashboards/README.md

@@ -74,8 +74,8 @@ percli apply -f perses/performance_statistics.yaml
 
 For detailed deployment instructions and platform-specific options, see:
 
-- **[Grafana Documentation](./grafana)** - JSON dashboards, operator usage, manual import
-- **[Perses Documentation](./perses)** - YAML specs, CLI usage, operator wrapping
+- **[Grafana Documentation](grafana)** - JSON dashboards, operator usage, manual import
+- **[Perses Documentation](perses)** - YAML specs, CLI usage, operator wrapping
 
 ## Contributing
 

examples/offline_inference/simple_profiling.py

@@ -2,46 +2,31 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import time
-import os
-
-os.environ["VLLM_USE_SPECIALIZED_MODELS"] = "1"
-os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "1"
 
 from vllm import LLM, SamplingParams
 
 # Sample prompts.
 prompts = [
-    [0] * 10_000,
-    [1] * 10_000,
-    [2] * 10_000,
-    [3] * 10_000,
-    [4] * 10_000,
-    [5] * 10_000,
-    [6] * 10_000,
-    [7] * 10_000,
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
 ]
 # Create a sampling params object.
-sampling_params = SamplingParams(temperature=0.0)
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 
 
 def main():
     # Create an LLM.
     llm = LLM(
-        model="nvidia/DeepSeek-V3.2-NVFP4",
-        tensor_parallel_size=4,
-        kernel_config={"enable_flashinfer_autotune": False},
+        model="facebook/opt-125m",
+        tensor_parallel_size=1,
         profiler_config={
             "profiler": "torch",
-            "torch_profiler_dir": f"./vllm_profile/bsz{len(prompts)}/",
+            "torch_profiler_dir": "./vllm_profile",
         },
-        enable_prefix_caching=False,
-        load_format="dummy",
-        compilation_config={"max_cudagraph_capture_size": 64},
-        speculative_config={"method": "mtp", "num_speculative_tokens": 3},
-        max_num_batched_tokens=32768,
     )
 
-    outputs = llm.generate(prompts, sampling_params)
     llm.start_profile()
 
     # Generate texts from the prompts. The output is a list of RequestOutput

tests/kernels/attention/test_cache.py

@@ -10,7 +10,7 @@ from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
 from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_dequantize
 from vllm.platforms import current_platform
-from vllm.utils.torch_utils import set_random_seed
+from vllm.utils.torch_utils import nvfp4_kv_cache_split_views, set_random_seed
 
 COPYING_DIRECTION = [("cuda", "cpu"), ("cuda", "cuda"), ("cpu", "cuda")]
 DTYPES = [torch.bfloat16, torch.float]
@@ -172,7 +172,7 @@ def test_reshape_and_cache(
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
-@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE)
+@pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPE + ["nvfp4"])
 @pytest.mark.parametrize("kv_cache_layout", CACHE_LAYOUTS)
 @pytest.mark.parametrize("kv_scale_type", KV_SCALE_TYPES)
 @pytest.mark.parametrize("implementation", RESHAPE_FLASH_IMPLEMENTATIONS)
@@ -202,6 +202,25 @@ def test_reshape_and_cache_flash(
     if kv_scale_type == "attn_head" and implementation != "cuda":
         pytest.skip("Only CUDA implementation supports attn_head scaling.")
 
+    if kv_cache_dtype == "nvfp4":
+        if not current_platform.has_device_capability(100):
+            pytest.skip("NVFP4 requires compute capability >= 10.0 (Blackwell).")
+        if implementation != "cuda":
+            pytest.skip("NVFP4 only supports CUDA implementation.")
+        if kv_scale_type != "tensor":
+            pytest.skip("NVFP4 only supports per-tensor scaling.")
+        if head_size % 16 != 0:
+            pytest.skip("NVFP4 requires head_size divisible by 16.")
+        if (head_size // 16) % 4 != 0:
+            pytest.skip(
+                "NVFP4 requires (head_size // 16) divisible by 4 "
+                "for 4x4 block scale swizzle."
+            )
+        if block_size % 4 != 0:
+            pytest.skip("NVFP4 requires block_size divisible by 4.")
+        if dtype not in (torch.float16, torch.bfloat16):
+            pytest.skip("NVFP4 quantization only supports fp16/bf16 input.")
+
     # fp8 conversion requires continugous memory buffer. Reduce the number of
     # blocks and tokens to consume less memory.
     num_tokens = num_tokens // 2
@@ -229,7 +248,23 @@ def test_reshape_and_cache_flash(
     del key_caches
     del value_caches
 
-    if kv_scale_type == "tensor":
+    # For nvfp4, the factory returns kv[:, 0] and kv[:, 1] like all dtypes.
+    # Split views are still needed for dequant verification.
+    key_scale_cache = None
+    value_scale_cache = None
+    nvfp4_key_data = None
+    nvfp4_value_data = None
+    if kv_cache_dtype == "nvfp4":
+        (nvfp4_key_data,), (key_scale_cache,) = nvfp4_kv_cache_split_views(key_cache)
+        (nvfp4_value_data,), (value_scale_cache,) = nvfp4_kv_cache_split_views(
+            value_cache
+        )
+
+    if kv_cache_dtype == "nvfp4":
+        # Global scale = amax / 448 (per-tensor)
+        k_scale = (key.abs().amax() / 448.0).to(torch.float32)
+        v_scale = (value.abs().amax() / 448.0).to(torch.float32)
+    elif kv_scale_type == "tensor":
         k_scale = (key.amax() / 64.0).to(torch.float32)
         v_scale = (value.amax() / 64.0).to(torch.float32)
     else:  # "attn_head"
@@ -240,8 +275,9 @@ def test_reshape_and_cache_flash(
         y = x if kv_cache_layout == "NHD" else x.permute(0, 2, 1, 3)
         return y.contiguous()
 
-    key_cache_compact = permute_and_compact(key_cache)
-    value_cache_compact = permute_and_compact(value_cache)
+    if kv_cache_dtype != "nvfp4":
+        key_cache_compact = permute_and_compact(key_cache)
+        value_cache_compact = permute_and_compact(value_cache)
 
     def convert_fp8_local(output, input, scale, kv_dtype):
         fp8_input = input.view(current_platform.fp8_dtype())
@@ -257,7 +293,7 @@ def test_reshape_and_cache_flash(
                 result = fp8_input.to(output.dtype) * scale.view(1, -1, 1, 1)
         output.copy_(result)
 
-    # Clone the KV caches.
+    # Clone the KV caches (for non-nvfp4, used as reference baseline).
     if kv_cache_dtype == "fp8":
         cloned_key_cache = torch.empty_like(key_cache_compact, dtype=torch.float16)
         convert_fp8_local(cloned_key_cache, key_cache_compact, k_scale, kv_cache_dtype)
@@ -265,25 +301,27 @@ def test_reshape_and_cache_flash(
         convert_fp8_local(
             cloned_value_cache, value_cache_compact, v_scale, kv_cache_dtype
         )
-    else:
+    elif kv_cache_dtype != "nvfp4":
         cloned_key_cache = key_cache_compact.clone()
         cloned_value_cache = value_cache_compact.clone()
+
     # Call the reshape_and_cache kernel.
     if implementation == "cuda":
-        opcheck(
-            torch.ops._C_cache_ops.reshape_and_cache_flash,
-            (
-                key,
-                value,
-                key_cache,
-                value_cache,
-                slot_mapping,
-                kv_cache_dtype,
-                k_scale,
-                v_scale,
-            ),
-            cond=(head_size == HEAD_SIZES[0]),
-        )
+        if kv_cache_dtype != "nvfp4":
+            opcheck(
+                torch.ops._C_cache_ops.reshape_and_cache_flash,
+                (
+                    key,
+                    value,
+                    key_cache,
+                    value_cache,
+                    slot_mapping,
+                    kv_cache_dtype,
+                    k_scale,
+                    v_scale,
+                ),
+                cond=(head_size == HEAD_SIZES[0]),
+            )
         ops.reshape_and_cache_flash(
             key,
             value,
@@ -309,6 +347,46 @@ def test_reshape_and_cache_flash(
             k_scale,
             v_scale,
         )
+
+    if kv_cache_dtype == "nvfp4":
+        # Verify NVFP4 by dequantizing the entire cache and comparing
+        # the written positions against original bf16 values.
+        # Same pattern as FP8: dequant whole cache, then extract and compare.
+        from tests.kernels.quantization.nvfp4_utils import (
+            dequant_nvfp4_kv_cache,
+        )
+
+        def dequant_nvfp4_cache_nhd(data_cache, scale_cache, global_scale):
+            # data_cache:  [N, T, H, data_dim]  NHD (contiguous inner dims)
+            # scale_cache: [N, T, H, scale_dim] NHD (contiguous inner dims)
+            # Permute to HND layout for the dequant utility.
+            data_hnd = data_cache.permute(0, 2, 1, 3)
+            scale_hnd = scale_cache.permute(0, 2, 1, 3)
+            result_hnd = dequant_nvfp4_kv_cache(
+                data_hnd, scale_hnd, global_scale, head_size, block_size
+            )
+            return result_hnd.permute(0, 2, 1, 3)  # back to [N, T, H, D]
+
+        result_key_cache = dequant_nvfp4_cache_nhd(
+            nvfp4_key_data, key_scale_cache, k_scale.item()
+        )
+        result_value_cache = dequant_nvfp4_cache_nhd(
+            nvfp4_value_data, value_scale_cache, v_scale.item()
+        )
+
+        # Flatten [num_blocks, block_size] → [num_slots] and index by slot_mapping.
+        num_slots = num_blocks * block_size
+        result_key_flat = result_key_cache.reshape(num_slots, num_heads, head_size)
+        result_value_flat = result_value_cache.reshape(num_slots, num_heads, head_size)
+
+        torch.testing.assert_close(
+            result_key_flat[slot_mapping], key.float(), atol=1.5, rtol=0.5
+        )
+        torch.testing.assert_close(
+            result_value_flat[slot_mapping], value.float(), atol=1.5, rtol=0.5
+        )
+        return
+
     key_cache_compact = permute_and_compact(key_cache)
     value_cache_compact = permute_and_compact(value_cache)
 

tests/kernels/quantization/nvfp4_utils.py

@@ -88,6 +88,60 @@ def break_fp4_bytes(a, dtype):
     return values.reshape(m, n * 2).to(dtype=dtype)
 
 
+def dequant_nvfp4_kv_cache(
+    fp4_data: torch.Tensor,
+    block_scale: torch.Tensor,
+    global_scale: float,
+    head_size: int,
+    block_size: int,
+) -> torch.Tensor:
+    """Dequantize an NVFP4 KV cache with 4x4-swizzled block scales.
+
+    The input must be in HND layout so that the last two dims are
+    (block_size, last_dim).  For NHD caches, permute to HND first.
+
+    Args:
+        fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4.
+        block_scale: [..., num_heads, block_size, head_size//16] fp8 block
+            scales (as uint8 or float8_e4m3fn).
+        global_scale: checkpoint dequant scale (k_scale or v_scale).
+        head_size: head dimension.
+        block_size: page size.
+
+    Returns:
+        [..., num_heads, block_size, head_size] float32.
+    """
+    data_dim = head_size // 2
+    scale_dim = head_size // 16
+
+    fp4_packed = fp4_data
+    sf_swizzled = block_scale.view(torch.uint8)
+
+    # Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
+    # [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
+    batch_shape = sf_swizzled.shape[:-2]
+    T, S = block_size, scale_dim
+    sg = S // 4
+    sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
+    ndim = sf_reshape.ndim
+    # Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
+    perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
+    sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
+    sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32)
+
+    # Unpack fp4
+    shape = fp4_packed.shape  # [..., T, data_dim]
+    fp4_flat = fp4_packed.reshape(-1, data_dim)
+    fp4_vals = break_fp4_bytes(fp4_flat, torch.float32)
+    fp4_vals = fp4_vals.reshape(*shape[:-1], head_size)
+
+    # Dequant: fp4_val * block_scale * global_scale per 16-element group
+    return (
+        fp4_vals.reshape(*shape[:-1], scale_dim, 16)
+        * (sf_f32 * global_scale).unsqueeze(-1)
+    ).reshape(*shape[:-1], head_size)
+
+
 def get_nvfp4_global_scale(a: torch.Tensor):
     return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
 

tests/models/language/generation/test_common.py

@@ -46,7 +46,7 @@ AITER_MODEL_LIST = [
         ),
         pytest.param(
             "openai-community/gpt2",  # gpt2
-            marks=[pytest.mark.core_model, pytest.mark.cpu_model],
+            marks=[pytest.mark.core_model],
         ),
         pytest.param("Milos/slovak-gpt-j-405M"),  # gptj
         pytest.param("bigcode/tiny_starcoder_py"),  # gpt_bigcode
@@ -143,11 +143,6 @@ def test_models(
         # in parts of the operators
         pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
 
-    if current_platform.is_cpu() and model in ("openai-community/gpt2",):
-        # These models are sensitive to the rounding error
-        # Fuse ops to reduce rounding
-        monkeypatch.setenv("VLLM_CPU_CI_ENV", "0")
-
     with hf_runner(model) as hf_model:
         hf_outputs = hf_model.generate_greedy_logprobs_limit(
             example_prompts, max_tokens, num_logprobs

tests/tool_parsers/test_qwen3xml_tool_parser.py

@@ -64,9 +64,6 @@ class TestQwen3xmlToolParser(ToolParserTests):
                 "test_empty_arguments": "Qwen3XML streaming has systematic issues",
                 "test_surrounding_text": "Qwen3XML streaming has systematic issues",
                 "test_escaped_strings": "Qwen3XML streaming has systematic issues",
-                "test_malformed_input": (
-                    "Qwen3XML parser is lenient with malformed input"
-                ),
                 "test_streaming_reconstruction": (
                     "Qwen3XML streaming reconstruction has known issues"
                 ),

tests/v1/attention/test_cutlass_fa3_sparse_backend.py

@@ -0,0 +1,968 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Backend integration tests for CUTLASS FA3 sparse MLA attention.
+
+Tests verify:
+  - Backend class properties
+  - Metadata builder (decode, prefill, mixed, topk clipping)
+  - KV cache write/read consistency
+  - Backend registration and selection
+"""
+
+import pytest
+import torch
+
+from vllm.v1.attention.ops.cutlass_fa3 import is_cutlass_fa3_available
+
+pytestmark = pytest.mark.skipif(
+    not is_cutlass_fa3_available(),
+    reason="CUTLASS FA3 not available (requires CUDA >= 12.4, SM90)",
+)
+
+
+# ─── TEST 2.1: Backend Class Properties ──────────────────────────────
+
+
+def test_backend_class_properties():
+    """Verify CutlassFA3MLASparseBackend class attributes."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseBackend,
+    )
+
+    assert CutlassFA3MLASparseBackend.get_name() == "CUTLASS_FA3_MLA_SPARSE"
+    assert CutlassFA3MLASparseBackend.is_mla() is True
+    assert CutlassFA3MLASparseBackend.is_sparse() is True
+    assert CutlassFA3MLASparseBackend.get_supported_head_sizes() == [576]
+    assert CutlassFA3MLASparseBackend.supported_kv_cache_dtypes == ["auto"]
+    assert CutlassFA3MLASparseBackend.get_supported_kernel_block_sizes() == [64]
+
+
+def test_backend_compute_capability():
+    """Verify SM90-only support."""
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseBackend,
+    )
+
+    assert CutlassFA3MLASparseBackend.supports_compute_capability(
+        DeviceCapability(major=9, minor=0)
+    )
+    assert not CutlassFA3MLASparseBackend.supports_compute_capability(
+        DeviceCapability(major=8, minor=0)
+    )
+    assert not CutlassFA3MLASparseBackend.supports_compute_capability(
+        DeviceCapability(major=10, minor=0)
+    )
+
+
+def test_backend_kv_cache_shape():
+    """Verify KV cache shape for BF16 format."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseBackend,
+    )
+
+    shape = CutlassFA3MLASparseBackend.get_kv_cache_shape(
+        num_blocks=100,
+        block_size=64,
+        num_kv_heads=1,
+        head_size=576,
+        cache_dtype_str="auto",
+    )
+    assert shape == (100, 64, 576)
+
+
+# ─── TEST 2.2: Backend Registration ──────────────────────────────────
+
+
+def test_backend_enum_registered():
+    """Verify CUTLASS_FA3_MLA_SPARSE is in the backend enum."""
+    from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+    assert hasattr(AttentionBackendEnum, "CUTLASS_FA3_MLA_SPARSE")
+    backend_enum = AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE
+    assert "cutlass_fa3_sparse" in backend_enum.get_path()
+
+
+def test_backend_class_loadable():
+    """Verify the backend class can be loaded from the enum."""
+    from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+    backend_cls = AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE.get_class()
+    assert backend_cls.get_name() == "CUTLASS_FA3_MLA_SPARSE"
+
+
+# ─── TEST 2.3: KV Cache Write/Read ───────────────────────────────────
+
+
+def test_kv_cache_write_read_consistency():
+    """Verify do_kv_cache_update writes match what forward_mqa would read."""
+    device = "cuda"
+    num_blocks = 4
+    block_size = 64
+    head_size = 576
+    kv_lora_rank = 512
+    qk_rope_head_dim = 64
+
+    # Create BF16 cache
+    cache = torch.zeros(
+        num_blocks, block_size, head_size, dtype=torch.bfloat16, device=device
+    )
+
+    # Write known values
+    T = 3
+    kv_c_normed = torch.randn(T, kv_lora_rank, dtype=torch.bfloat16, device=device)
+    k_pe = torch.randn(T, 1, qk_rope_head_dim, dtype=torch.bfloat16, device=device)
+    slot_mapping = torch.tensor([0, 1, 2], dtype=torch.int64, device=device)
+    k_scale = torch.ones(1, dtype=torch.float32, device=device)
+
+    from vllm import _custom_ops as ops
+
+    ops.concat_and_cache_mla(
+        kv_c_normed,
+        k_pe.squeeze(1),
+        cache,
+        slot_mapping,
+        kv_cache_dtype="auto",
+        scale=k_scale,
+    )
+
+    # Read back via flatten + split (same as forward_mqa does)
+    S = num_blocks * block_size
+    kv_flat = cache.reshape(S, head_size)
+    c_kv_read = kv_flat[:T, :kv_lora_rank]
+    k_rope_read = kv_flat[:T, kv_lora_rank:]
+
+    # Verify consistency
+    torch.testing.assert_close(c_kv_read, kv_c_normed, rtol=1e-3, atol=1e-3)
+    torch.testing.assert_close(k_rope_read, k_pe.squeeze(1), rtol=1e-3, atol=1e-3)
+
+
+def test_kv_cache_dtype_auto():
+    """Verify kv_cache_dtype='auto' uses BF16 direct copy."""
+    device = "cuda"
+    cache = torch.zeros(1, 64, 576, dtype=torch.bfloat16, device=device)
+
+    kv_c = torch.randn(1, 512, dtype=torch.bfloat16, device=device)
+    k_pe = torch.randn(1, 1, 64, dtype=torch.bfloat16, device=device)
+    slot_mapping = torch.tensor([0], dtype=torch.int64, device=device)
+    k_scale = torch.ones(1, dtype=torch.float32, device=device)
+
+    from vllm import _custom_ops as ops
+
+    ops.concat_and_cache_mla(
+        kv_c, k_pe.squeeze(1), cache, slot_mapping, kv_cache_dtype="auto", scale=k_scale
+    )
+
+    assert cache.dtype == torch.bfloat16
+
+
+# ─── TEST 2.4: Edge Cases ────────────────────────────────────────────
+
+
+def test_empty_kv_cache():
+    """Verify do_kv_cache_update handles empty cache gracefully."""
+
+    kv_cache = torch.empty(0, device="cuda")
+    # Should return without error (numel() == 0 check)
+    # We call the static method from parent class directly
+    from vllm.v1.attention.backend import SparseMLAAttentionImpl
+
+    SparseMLAAttentionImpl.do_kv_cache_update(
+        None,
+        kv_c_normed=torch.empty(0),
+        k_pe=torch.empty(0),
+        kv_cache=kv_cache,
+        slot_mapping=torch.empty(0),
+        kv_cache_dtype="auto",
+        k_scale=torch.ones(1),
+    )
+
+
+# ─── TEST 2.5: Valid Counts from Index Conversion ───────────────────
+
+
+def test_triton_convert_valid_counts():
+    """Verify triton_convert_req_index_to_global_index with return_valid_counts.
+
+    This tests the core fix mechanism: the Triton kernel atomically counts
+    valid (non -1) entries per row while converting indices.
+    """
+    from vllm.v1.attention.backends.mla.sparse_utils import (
+        triton_convert_req_index_to_global_index,
+    )
+
+    device = "cuda"
+    T = 4
+    topk = 128
+    num_blocks = 16
+    block_size = 64
+
+    req_id = torch.zeros(T, dtype=torch.int32, device=device)
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
+        0
+    )  # [1, num_blocks]
+
+    # Create topk_indices with varying valid entries per token
+    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
+    expected_valid = [1, 10, 50, 100]
+    for i in range(T):
+        nv = expected_valid[i]
+        # Use indices within the valid range
+        topk_indices[i, :nv] = torch.randint(
+            0,
+            num_blocks * block_size,
+            (nv,),
+            dtype=torch.int32,
+            device=device,
+        )
+
+    global_idx, valid_counts = triton_convert_req_index_to_global_index(
+        req_id,
+        block_table,
+        topk_indices,
+        BLOCK_SIZE=block_size,
+        NUM_TOPK_TOKENS=topk,
+        return_valid_counts=True,
+    )
+
+    # Verify valid counts match expected
+    for i in range(T):
+        assert valid_counts[i].item() == expected_valid[i], (
+            f"Token {i}: expected {expected_valid[i]} valid, "
+            f"got {valid_counts[i].item()}"
+        )
+
+    # Verify -1 propagation
+    for i in range(T):
+        nv = expected_valid[i]
+        # Entries beyond valid should be -1
+        assert (global_idx[i, nv:] == -1).all(), (
+            f"Token {i}: entries beyond valid count should be -1"
+        )
+
+
+# ─── TEST 2.6: Prefill Metadata Correctness ─────────────────────────
+
+
+def test_prefill_cache_seqlens_vs_valid_counts():
+    """Verify metadata cache_seqlens = min(seq_len, topk) and that the
+    forward_mqa fix overrides with valid_counts.
+
+    The metadata builder computes cache_seqlens as min(seq_len, topk).
+    For prefill tokens, this can exceed the actual valid topk entries.
+    The fix in forward_mqa uses valid_counts instead.
+    """
+    import numpy as np
+
+    device = "cuda"
+
+    # Simulate a prefill batch: 1 request, 4 tokens, seq_len=4
+    num_reqs = 1
+    T = 4
+    topk = 2048
+    seq_len = 4
+
+    # The metadata builder's logic (simplified):
+    starts = np.array([0, T], dtype=np.int32)
+    seg_lens = np.diff(starts)  # [4]
+    seq_lens_np = np.array([seq_len], dtype=np.int32)
+
+    per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), topk)  # [4, 4, 4, 4]
+
+    # This is what the metadata builder produces:
+    assert all(per_tok_seqlens == 4), (
+        "Metadata cache_seqlens should be min(seq_len, topk) = 4"
+    )
+
+    # But the actual valid entries per token (with causal masking):
+    # Token 0: 1 valid entry, Token 1: 2, Token 2: 3, Token 3: 4
+    expected_valid = [1, 2, 3, 4]
+
+    # The fix in forward_mqa computes valid_counts from the page_table
+    # and uses those as cache_seqlens. Verify the fix produces correct
+    # valid counts:
+    from vllm.v1.attention.backends.mla.sparse_utils import (
+        triton_convert_req_index_to_global_index,
+    )
+
+    req_id = torch.zeros(T, dtype=torch.int32, device=device)
+    block_table = torch.arange(32, dtype=torch.int32, device=device).unsqueeze(0)
+
+    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
+    for i in range(T):
+        nv = expected_valid[i]
+        topk_indices[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device)
+
+    _, valid_counts = triton_convert_req_index_to_global_index(
+        req_id,
+        block_table,
+        topk_indices,
+        BLOCK_SIZE=64,
+        NUM_TOPK_TOKENS=topk,
+        return_valid_counts=True,
+    )
+
+    for i in range(T):
+        assert valid_counts[i].item() == expected_valid[i], (
+            f"Token {i}: valid_counts should be {expected_valid[i]}, "
+            f"got {valid_counts[i].item()}"
+        )
+
+
+# ─── TEST 2.7: Clamp -1 to 0 Safety ─────────────────────────────────
+
+
+def test_global_idx_clamp_safety():
+    """Verify clamping -1 page indices to 0 prevents OOB access."""
+    device = "cuda"
+
+    # Create a page_table with -1 entries
+    page_table = torch.tensor(
+        [[5, 10, -1, -1], [3, -1, -1, -1]],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    # Clamp -1 to 0
+    clamped = page_table.clamp(min=0)
+
+    # Verify
+    expected = torch.tensor(
+        [[5, 10, 0, 0], [3, 0, 0, 0]],
+        dtype=torch.int32,
+        device=device,
+    )
+    assert torch.equal(clamped, expected), (
+        f"Clamped page_table doesn't match expected: {clamped} vs {expected}"
+    )
+
+
+# ─── TEST 2.8: In-place clamp correctness ───────────────────────────
+
+
+def test_inplace_clamp_no_negative_indices():
+    """Verify in-place clamp_(min=0) on global_idx leaves no -1 entries.
+
+    The review-fixed code uses clamp_() (in-place) instead of clamp()
+    to avoid unnecessary tensor allocations during CUDA graph capture.
+    """
+    device = "cuda"
+
+    # Create a global_idx tensor with -1 entries
+    global_idx = torch.tensor(
+        [[100, 200, -1, -1, -1], [50, -1, -1, -1, -1]],
+        dtype=torch.int32,
+        device=device,
+    )
+
+    # In-place clamp
+    global_idx.clamp_(min=0)
+
+    # Verify no -1 entries remain
+    assert (global_idx >= 0).all(), (
+        f"In-place clamp should remove all -1 entries: {global_idx}"
+    )
+    # Verify valid entries are preserved
+    assert global_idx[0, 0].item() == 100
+    assert global_idx[0, 1].item() == 200
+    assert global_idx[1, 0].item() == 50
+
+
+# ─── TEST 2.9: Full fix flow with index conversion ──────────────────
+
+
+def test_full_fix_flow_valid_counts_and_clamp():
+    """End-to-end test of the complete fix flow:
+    1. triton_convert_req_index_to_global_index with return_valid_counts=True
+    2. In-place clamp global_idx to replace -1 with 0
+    3. In-place clamp valid_counts to min=1
+    4. Use valid_counts as cache_seqlens
+
+    This simulates what forward_mqa does after the fix.
+    """
+    from vllm.v1.attention.backends.mla.sparse_utils import (
+        triton_convert_req_index_to_global_index,
+    )
+
+    device = "cuda"
+    T = 4
+    topk = 128
+    num_blocks = 16
+    block_size = 64
+
+    req_id = torch.zeros(T, dtype=torch.int32, device=device)
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
+        0
+    )
+
+    # Simulate causal prefill: token i has (i+1) valid entries
+    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
+    expected_valid = [1, 2, 3, 4]
+    for i in range(T):
+        nv = expected_valid[i]
+        topk_indices[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device)
+
+    # Step 1: Convert with valid counts
+    global_idx, valid_counts = triton_convert_req_index_to_global_index(
+        req_id,
+        block_table,
+        topk_indices,
+        BLOCK_SIZE=block_size,
+        NUM_TOPK_TOKENS=topk,
+        return_valid_counts=True,
+    )
+
+    # Step 2: In-place clamp global_idx (no -1 entries after)
+    global_idx.clamp_(min=0)
+    assert (global_idx >= 0).all(), "No -1 entries should remain after clamp_"
+
+    # Step 3: In-place clamp valid_counts to min=1
+    valid_counts.clamp_(min=1)
+    cache_seqlens = valid_counts
+
+    # Step 4: Verify valid counts match expected
+    for i in range(T):
+        assert cache_seqlens[i].item() == expected_valid[i], (
+            f"Token {i}: expected cache_seqlens={expected_valid[i]}, "
+            f"got {cache_seqlens[i].item()}"
+        )
+
+    # Step 5: Verify that for each token, entries 0..cache_seqlens-1 in
+    # global_idx are valid (non-zero, since we clamped -1 to 0 for the
+    # entries beyond valid_counts, the valid entries at positions 0..nv-1
+    # should be the actual converted indices)
+    for i in range(T):
+        nv = expected_valid[i]
+        valid_region = global_idx[i, :nv]
+        # Valid region should have specific converted values from block_table
+        # (not just zeros from clamping)
+        # For indices [0, 1, ..., nv-1] with block_size=64:
+        #   block_id = index // 64, inblock_off = index % 64
+        #   out = block_table[0, block_id] * 64 + inblock_off
+        for j in range(nv):
+            block_id = j // block_size
+            inblock_off = j % block_size
+            expected_val = block_table[0, block_id].item() * block_size + inblock_off
+            assert valid_region[j].item() == expected_val, (
+                f"Token {i}, position {j}: expected {expected_val}, "
+                f"got {valid_region[j].item()}"
+            )
+
+
+# ─── TEST 2.10: CUDA Graph Padding Fix ─────────────────────────────
+# These tests verify the fix for Issue 2: RuntimeError when
+# num_actual_tokens (padded) != sum(seg_lens) (real tokens).
+# This is the core bug that caused the crash during lm_eval with
+# 32 concurrent requests on DeepSeek-V3.2.
+
+
+def _make_mock_vllm_config(max_tokens=512):
+    """Create a mock VllmConfig for metadata builder tests."""
+    from unittest.mock import MagicMock
+
+    vllm_config = MagicMock()
+    vllm_config.scheduler_config.max_num_batched_tokens = max_tokens
+    vllm_config.speculative_config = None
+    vllm_config.parallel_config.decode_context_parallel_size = 1
+    return vllm_config
+
+
+def test_metadata_builder_cuda_graph_padding():
+    """Verify build() handles CUDA graph padding (T > actual_tokens).
+
+    Reproduces the exact crash from Issue 2:
+      RuntimeError: The size of tensor a (32) must match the size
+      of tensor b (31) at non-singleton dimension 0
+
+    This happens when num_actual_tokens=32 (padded for CUDA graph)
+    but only 31 real tokens exist (one request completed mid-batch).
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+    topk = 2048
+
+    # Mock kv_cache_spec
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    # Mock vllm_config
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = topk
+
+    # Simulate the crash scenario: 31 real tokens padded to 32
+    padded_T = 32
+    real_tokens = 31
+    num_reqs_padded = 32  # padded request count
+
+    # Accurately mock gpu_model_runner.py's padding behavior:
+    # query_start_loc.cpu[:num_reqs_padded+1] = [:33], 33 entries
+    # Real entries: [0,1,...,31], Padding: [31] (repeats last value)
+    query_start_loc_cpu = list(range(real_tokens + 1)) + [real_tokens]
+    # seq_lens_cpu[:num_reqs_padded] = [:32], 32 entries
+    # Real entries: [100]*31, Padding: [0] (stale/zero for padding slot)
+    seq_lens_cpu = [100] * real_tokens + [0]
+
+    # Build the mock CommonAttentionMetadata
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T  # PADDED to 32
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = num_reqs_padded  # gpu_model_runner passes padded count
+    cm.max_query_len = 1
+    cm.max_seq_len = 100
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(
+        num_reqs_padded, 4, dtype=torch.int32, device=device
+    )
+
+    # This should NOT raise RuntimeError
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify metadata shapes match padded T
+    assert metadata.req_id_per_token.shape[0] == padded_T, (
+        f"req_id_per_token should have padded size {padded_T}, "
+        f"got {metadata.req_id_per_token.shape[0]}"
+    )
+    assert metadata.cache_seqlens.shape[0] == padded_T, (
+        f"cache_seqlens should have padded size {padded_T}, "
+        f"got {metadata.cache_seqlens.shape[0]}"
+    )
+    assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
+    assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
+
+    # Verify real data portion is correct
+    for i in range(real_tokens):
+        assert metadata.req_id_per_token[i].item() == i, (
+            f"Token {i}: req_id should be {i}, "
+            f"got {metadata.req_id_per_token[i].item()}"
+        )
+        assert metadata.cache_seqlens[i].item() == 100, (
+            f"Token {i}: cache_seqlens should be 100, "
+            f"got {metadata.cache_seqlens[i].item()}"
+        )
+
+    # Verify padding tokens have safe defaults
+    assert metadata.req_id_per_token[real_tokens].item() == 0, (
+        "Padding token req_id should be 0"
+    )
+    assert metadata.cache_seqlens[real_tokens].item() >= 1, (
+        "Padding token cache_seqlens should be >= 1 (safe minimum)"
+    )
+
+    # Verify cu_seqlens_q is [0, 1, 2, ..., padded_T] (always correct)
+    for i in range(padded_T + 1):
+        assert metadata.cu_seqlens_q[i].item() == i, (
+            f"cu_seqlens_q[{i}] should be {i}, got {metadata.cu_seqlens_q[i].item()}"
+        )
+
+    # Verify cu_seqlens_k is monotonically non-decreasing
+    for i in range(padded_T):
+        assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item(), (
+            f"cu_seqlens_k must be non-decreasing at index {i}: "
+            f"{metadata.cu_seqlens_k[i].item()} -> {metadata.cu_seqlens_k[i + 1].item()}"
+        )
+    # Verify cu_seqlens_k at the real/padding boundary
+    assert metadata.cu_seqlens_k[real_tokens].item() == real_tokens * 100, (
+        f"cu_seqlens_k[{real_tokens}] should be {real_tokens * 100}, "
+        f"got {metadata.cu_seqlens_k[real_tokens].item()}"
+    )
+
+
+@pytest.mark.parametrize(
+    "real_tokens,padded_T",
+    [
+        (1, 2),  # minimal padding
+        (3, 32),  # large padding gap
+        (7, 8),  # small batch
+        (15, 16),  # medium batch
+        (31, 32),  # the exact crash scenario
+        (100, 104),  # larger padding gap
+    ],
+)
+def test_metadata_builder_cuda_graph_padding_various(real_tokens, padded_T):
+    """Verify build() handles various CUDA graph padding scenarios.
+
+    Uses accurate mock that matches gpu_model_runner.py's padding behavior:
+    - query_start_loc_cpu has num_reqs_padded+1 entries (with padded suffix)
+    - seq_lens_cpu has num_reqs_padded entries (with stale padding entries)
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = max(512, padded_T + 1)  # ensure buffer large enough
+    block_size = 64
+    topk = 2048
+    num_reqs_padded = padded_T  # For decode-only, padded_T == num_reqs_padded
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = topk
+
+    # Accurate mock: query_start_loc_cpu[:num_reqs_padded+1]
+    # Real entries [0,1,...,real_tokens], then (num_reqs_padded - real_tokens)
+    # padding entries all equal to real_tokens (flat, non-decreasing)
+    query_start_loc_cpu = list(range(real_tokens + 1))
+    num_padding_reqs = num_reqs_padded - real_tokens
+    query_start_loc_cpu += [real_tokens] * num_padding_reqs
+
+    # seq_lens_cpu[:num_reqs_padded] — padding entries are stale (zero)
+    seq_lens_cpu = [200] * real_tokens + [0] * num_padding_reqs
+
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = num_reqs_padded
+    cm.max_query_len = 1
+    cm.max_seq_len = 200
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(
+        max(num_reqs_padded, 1), 4, dtype=torch.int32, device=device
+    )
+
+    # Should NOT raise any errors
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify shapes match padded T
+    assert metadata.req_id_per_token.shape[0] == padded_T
+    assert metadata.cache_seqlens.shape[0] == padded_T
+    assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
+    assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
+    assert metadata.num_actual_tokens == padded_T
+
+    # Verify real portion
+    for i in range(real_tokens):
+        assert metadata.req_id_per_token[i].item() == i
+        assert metadata.cache_seqlens[i].item() == 200
+
+    # Verify padding
+    for i in range(real_tokens, padded_T):
+        assert metadata.req_id_per_token[i].item() == 0
+        assert metadata.cache_seqlens[i].item() >= 1
+
+    # Verify cu_seqlens_q is [0, 1, ..., padded_T]
+    for i in range(padded_T + 1):
+        assert metadata.cu_seqlens_q[i].item() == i
+
+    # Verify cu_seqlens_k monotonicity
+    for i in range(padded_T):
+        assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item()
+    # Verify cu_seqlens_k at boundary
+    assert metadata.cu_seqlens_k[real_tokens].item() == real_tokens * 200
+
+
+def test_metadata_builder_no_padding():
+    """Verify build() still works correctly when T == actual_tokens (no padding)."""
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = 2048
+
+    # No padding: T == real tokens
+    T = 4
+    query_start_loc_cpu = [0, 1, 2, 3, 4]  # 4 decode tokens
+    seq_lens_cpu = [50, 100, 150, 200]
+
+    cm = MagicMock()
+    cm.num_actual_tokens = T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = 4
+    cm.max_query_len = 1
+    cm.max_seq_len = 200
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(4, 4, dtype=torch.int32, device=device)
+
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    assert metadata.req_id_per_token.shape[0] == T
+    assert metadata.cache_seqlens.shape[0] == T
+    assert metadata.num_actual_tokens == T
+
+    # Verify exact values
+    assert metadata.req_id_per_token[0].item() == 0
+    assert metadata.req_id_per_token[1].item() == 1
+    assert metadata.req_id_per_token[2].item() == 2
+    assert metadata.req_id_per_token[3].item() == 3
+    assert metadata.cache_seqlens[0].item() == 50
+    assert metadata.cache_seqlens[1].item() == 100
+    assert metadata.cache_seqlens[2].item() == 150
+    assert metadata.cache_seqlens[3].item() == 200
+
+
+def test_metadata_builder_mixed_prefill_decode_with_padding():
+    """Verify build() handles mixed prefill+decode with CUDA graph padding.
+
+    This tests a more complex scenario: 2 decode tokens + 3 prefill tokens
+    from 3 requests, padded from 5 to 8 tokens.
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = 2048
+
+    # 3 real requests: req0 (1 decode token), req1 (1 decode token),
+    #                  req2 (3 prefill tokens)
+    # Total: 5 real tokens, padded to 8 tokens, 8 padded request slots
+    real_tokens = 5
+    num_real_reqs = 3
+    padded_T = 8
+    num_reqs_padded = 8  # padded request count
+
+    # Accurate: query_start_loc_cpu[:num_reqs_padded+1] = 9 entries
+    # Real: [0, 1, 2, 5], Padding: [5, 5, 5, 5, 5]
+    query_start_loc_cpu = [0, 1, 2, 5] + [5] * (num_reqs_padded - num_real_reqs)
+    # seq_lens_cpu[:num_reqs_padded] = 8 entries
+    seq_lens_cpu = [100, 200, 3] + [0] * (num_reqs_padded - num_real_reqs)
+
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = num_reqs_padded
+    cm.max_query_len = 3
+    cm.max_seq_len = 200
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(
+        num_reqs_padded, 4, dtype=torch.int32, device=device
+    )
+
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify shapes
+    assert metadata.req_id_per_token.shape[0] == padded_T
+    assert metadata.cache_seqlens.shape[0] == padded_T
+
+    # Verify req_id mapping
+    assert metadata.req_id_per_token[0].item() == 0  # req0, decode
+    assert metadata.req_id_per_token[1].item() == 1  # req1, decode
+    assert metadata.req_id_per_token[2].item() == 2  # req2, prefill tok0
+    assert metadata.req_id_per_token[3].item() == 2  # req2, prefill tok1
+    assert metadata.req_id_per_token[4].item() == 2  # req2, prefill tok2
+    # Padding tokens
+    assert metadata.req_id_per_token[5].item() == 0
+    assert metadata.req_id_per_token[6].item() == 0
+    assert metadata.req_id_per_token[7].item() == 0
+
+    # Verify cache_seqlens
+    assert metadata.cache_seqlens[0].item() == 100  # req0 seq_len
+    assert metadata.cache_seqlens[1].item() == 200  # req1 seq_len
+    assert metadata.cache_seqlens[2].item() == 3  # req2 seq_len
+    assert metadata.cache_seqlens[3].item() == 3  # req2 seq_len
+    assert metadata.cache_seqlens[4].item() == 3  # req2 seq_len
+    # Padding (default = 1)
+    assert metadata.cache_seqlens[5].item() >= 1
+    assert metadata.cache_seqlens[6].item() >= 1
+    assert metadata.cache_seqlens[7].item() >= 1
+
+
+# ─── TEST 2.11: Zero Real Tokens Edge Case (Review Issue #3) ─────
+# Tests the edge case where ALL tokens are padding (actual_tokens=0).
+# This can happen during CUDA graph warmup/capture with dummy batches.
+
+
+def test_metadata_builder_zero_real_tokens():
+    """Verify build() handles the case where all tokens are padding.
+
+    This edge case can occur during CUDA graph warmup or capture where
+    dummy batches may have zero real tokens but T > 0 (padded size).
+    """
+    from unittest.mock import MagicMock
+
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseMetadataBuilder,
+    )
+
+    device = "cuda"
+    max_tokens = 512
+    block_size = 64
+
+    kv_cache_spec = MagicMock()
+    kv_cache_spec.block_size = block_size
+
+    vllm_config = _make_mock_vllm_config(max_tokens)
+
+    builder = CutlassFA3MLASparseMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layers.0.self_attn"],
+        vllm_config=vllm_config,
+        device=torch.device(device),
+    )
+    builder.topk_tokens = 2048
+
+    # Zero real tokens, padded to 4
+    # This happens when query_start_loc = [0] only (1 entry, no requests)
+    # and num_actual_tokens is still the padded count.
+    padded_T = 4
+    real_tokens = 0
+    # query_start_loc_cpu with a single entry means 0 requests
+    query_start_loc_cpu = [0]
+    seq_lens_cpu = []
+
+    cm = MagicMock()
+    cm.num_actual_tokens = padded_T
+    cm.query_start_loc_cpu = query_start_loc_cpu
+    cm.seq_lens_cpu = seq_lens_cpu
+    cm.num_reqs = 0
+    cm.max_query_len = 0
+    cm.max_seq_len = 0
+    cm.query_start_loc = torch.tensor(
+        query_start_loc_cpu, dtype=torch.int32, device=device
+    )
+    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
+    cm.block_table_tensor = torch.zeros(1, 4, dtype=torch.int32, device=device)
+
+    # Should NOT raise any errors
+    metadata = builder.build(
+        common_prefix_len=0,
+        common_attn_metadata=cm,
+    )
+
+    # Verify shapes match padded T
+    assert metadata.req_id_per_token.shape[0] == padded_T
+    assert metadata.cache_seqlens.shape[0] == padded_T
+    assert metadata.cu_seqlens_q.shape[0] == padded_T + 1
+    assert metadata.cu_seqlens_k.shape[0] == padded_T + 1
+
+    # All tokens are padding — verify safe defaults
+    for i in range(padded_T):
+        assert metadata.req_id_per_token[i].item() == 0
+        assert metadata.cache_seqlens[i].item() >= 1
+
+    # cu_seqlens_k should be monotonically non-decreasing
+    for i in range(padded_T):
+        assert metadata.cu_seqlens_k[i + 1].item() >= metadata.cu_seqlens_k[i].item()
+
+
+# ─── TEST 2.12: Batch Size Gating Constant ───────────────────────────
+
+
+def test_batch_size_gating_threshold():
+    """Verify MAX_BATCH_SIZE_FOR_FA3 is 16 and controls routing."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        MAX_BATCH_SIZE_FOR_FA3,
+        _flashmla_sparse_available,
+    )
+
+    assert MAX_BATCH_SIZE_FOR_FA3 == 16
+    # On SM90 builds, FlashMLA fallback should be available
+    # (unless FlashMLA was explicitly excluded from the build)
+    assert isinstance(_flashmla_sparse_available, bool)
+
+
+# ─── TEST 2.13: FlashMLA Fallback Head Padding ──────────────────────
+
+
+def test_flashmla_fallback_head_padding():
+    """Verify FlashMLA fallback head padding constant is 64 for SM90."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        _FLASHMLA_SM90_HEAD_PADDING,
+    )
+
+    assert _FLASHMLA_SM90_HEAD_PADDING == 64, (
+        f"SM90 head padding should be 64, got {_FLASHMLA_SM90_HEAD_PADDING}"
+    )
+
+
+# ─── TEST 2.14: Forward MQA Dispatch Verification ────────────────────
+
+
+def test_forward_mqa_has_fa3_and_fallback_methods():
+    """Verify CutlassFA3MLASparseImpl has both kernel dispatch methods."""
+    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
+        CutlassFA3MLASparseImpl,
+    )
+
+    assert hasattr(CutlassFA3MLASparseImpl, "_forward_fa3"), (
+        "CutlassFA3MLASparseImpl should have _forward_fa3 method"
+    )
+    assert hasattr(CutlassFA3MLASparseImpl, "_forward_flashmla_bf16_fallback"), (
+        "CutlassFA3MLASparseImpl should have _forward_flashmla_bf16_fallback method"
+    )

vllm/config/cache.py

@@ -30,6 +30,7 @@ CacheDType = Literal[
     "turboquant_3bit_nc",
     "int8_per_token_head",
     "fp8_per_token_head",
+    "nvfp4",
 ]
 MambaDType = Literal["auto", "float32", "float16"]
 MambaCacheMode = Literal["all", "align", "none"]
@@ -101,8 +102,6 @@ class CacheConfig:
     kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)
     """Layer patterns to skip KV cache quantization. Accepts layer indices
     (e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""
-    cpu_kvcache_space_bytes: int | None = None
-    """(CPU backend only) CPU key-value cache space."""
     mamba_page_size_padded: int | None = None
     """ Optional override for mamba page size; used by hybrid mamba/attention
     models to ensure exact alignment with attention page size."""
@@ -183,7 +182,6 @@ class CacheConfig:
             "num_gpu_blocks_override",
             "enable_prefix_caching",
             "prefix_caching_hash_algo",
-            "cpu_kvcache_space_bytes",
             "mamba_page_size_padded",
             "user_specified_block_size",
             "user_specified_mamba_block_size",

vllm/config/compilation.py

@@ -737,8 +737,6 @@ class CompilationConfig:
         "vllm::kda_attention",
         "vllm::sparse_attn_indexer",
         "vllm::rocm_aiter_sparse_attn_indexer",
-        # For specialized models
-        "vllm::monolithic_attn",
     ]
 
     def compute_hash(self) -> str:

vllm/entrypoints/pooling/pooling/protocol.py

@@ -97,6 +97,11 @@ class IOProcessorRequest(PoolingBasicRequestMixin, EncodingRequestMixin, Generic
             max_total_tokens_param="max_model_len",
         )
 
+    def to_pooling_params(self):
+        return PoolingParams(
+            task=self.task,
+        )
+
 
 class IOProcessorResponse(OpenAIBaseModel, Generic[T]):
     request_id: str | None = None

vllm/envs.py

@@ -216,7 +216,6 @@ if TYPE_CHECKING:
     VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False
     VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
     VLLM_ALLREDUCE_USE_FLASHINFER: bool = False
-    VLLM_USE_SPECIALIZED_MODELS: bool = False
     VLLM_TUNED_CONFIG_FOLDER: str | None = None
     VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
     VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False
@@ -1521,10 +1520,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool(
         int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0"))
     ),
-    # Whether to enable specialized model implementations when available.
-    "VLLM_USE_SPECIALIZED_MODELS": lambda: bool(
-        int(os.getenv("VLLM_USE_SPECIALIZED_MODELS", "0"))
-    ),
     # Experimental: use this to enable MCP tool calling for non harmony models
     "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool(
         int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0"))

vllm/model_executor/kernels/linear/__init__.py

@@ -186,12 +186,13 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
 
 # in priority/performance order (when available)
 _POSSIBLE_FP8_BLOCK_KERNELS: dict[
-    PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel]]
+    PlatformEnum, list[type[Fp8BlockScaledMMLinearKernel | FP8ScaledMMLinearKernel]]
 ] = {
     PlatformEnum.CUDA: [
         FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
         DeepGemmFp8BlockScaledMMKernel,
         CutlassFp8BlockScaledMMKernel,
+        MarlinFP8ScaledMMLinearKernel,
         TritonFp8BlockScaledMMKernel,
     ],
     PlatformEnum.ROCM: [
@@ -392,6 +393,19 @@ def init_fp8_linear_kernel(
                 scope="global",
             )
 
+        # TODO make scaled_mm kernels inherit from MMLinearKernel
+        # only MarlinFP8ScaledMMLinearKernel is a type of FP8ScaledMMLinearKernel.
+        if issubclass(kernel_type, FP8ScaledMMLinearKernel):
+            return kernel_type(
+                scaled_mm_linear_kernel_config,
+                layer_param_names=[
+                    "weight",
+                    "weight_scale",
+                    "input_scale",
+                    "input_scale_ub",
+                ],
+            )
+
         return kernel_type(
             scaled_mm_linear_kernel_config,
         )
@@ -399,7 +413,7 @@ def init_fp8_linear_kernel(
     else:
         kernel_type = choose_scaled_mm_linear_kernel(
             config=scaled_mm_linear_kernel_config,
-            possible_kernels=_POSSIBLE_FP8_KERNELS,  # type: ignore[misc]
+            possible_kernels=_POSSIBLE_FP8_KERNELS,  # type: ignore[arg-type]
             force_kernel=force_kernel,
         )
         if module_name:

vllm/model_executor/layers/attention/attention.py

@@ -387,7 +387,9 @@ class Attention(nn.Module, AttentionLayerBase):
         self.query_quant = None
         if (
             self.impl.supports_quant_query_input
-            and self.kv_cache_dtype.startswith("fp8")
+            and (
+                self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4"
+            )
             and not self.kv_cache_dtype.endswith("per_token_head")
         ):
             is_per_head = (
@@ -492,7 +494,7 @@ class Attention(nn.Module, AttentionLayerBase):
             # which reduces overheads during decoding.
             # Otherwise queries are quantized using custom ops
             # which causes decoding overheads
-            assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"}
+            assert self.kv_cache_dtype in {"fp8", "fp8_e4m3", "nvfp4"}
 
             # check if query quantization is supported
             if self.impl.supports_quant_query_input:

vllm/model_executor/layers/attention/mla_attention.py

@@ -367,6 +367,16 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 "KV cache format, please set `--attention-backend FLASHMLA_SPARSE`"
             )
 
+        # CUTLASS FA3 MLA Sparse requires BF16 KV cache — force "auto" dtype
+        if self.attn_backend.get_name() == "CUTLASS_FA3_MLA_SPARSE":
+            if cache_config is not None:
+                cache_config.cache_dtype = "auto"
+            kv_cache_dtype = "auto"
+            logger.info_once(
+                "CUTLASS FA3 MLA Sparse backend requires BF16 KV cache. "
+                "Setting kv_cache_dtype to 'auto' (BF16)."
+            )
+
         # Initialize KV cache quantization attributes
         self.kv_cache_dtype = kv_cache_dtype
         self.calculate_kv_scales = calculate_kv_scales

vllm/model_executor/model_loader/dummy_loader.py

@@ -43,13 +43,6 @@ class DummyModelLoader(BaseModelLoader):
                 # random values to the weights.
                 initialize_dummy_weights(layer, model_config)
 
-                # Some models build derived weights from loaded parameters instead of
-                # storing them in checkpoints. Rebuild those tensors for dummy load.
-                for layer in model.modules():
-                    fuse_indexer_weights = getattr(layer, "fuse_indexer_weights", None)
-                    if callable(fuse_indexer_weights):
-                        fuse_indexer_weights()
-
     def _process_online_quant_layer(
         self,
         layer: nn.Module,

vllm/model_executor/models/deepseek_mtp.py

@@ -30,6 +30,7 @@ from .deepseek_v2 import (
     DeepseekV2DecoderLayer,
     DeepseekV2MixtureOfExperts,
     DeepseekV2MoE,
+    _try_load_fp8_indexer_wk,
     get_spec_layer_idx_from_weight_name,
 )
 from .utils import maybe_prefix
@@ -190,10 +191,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
         )
         # Set MoE hyperparameters
         self.set_moe_parameters()
-        self.is_fp4_ckpt = (
-            self.quant_config is not None
-            and self.quant_config.get_name() == "modelopt_fp4"
-        )
 
     def set_moe_parameters(self):
         self.expert_weights = []
@@ -248,13 +245,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
             ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
         ]
 
-        if self.is_fp4_ckpt:
-            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
-            indexer_fused_mapping = [
-                ("wk_weights_proj", "wk", 0),
-                ("wk_weights_proj", "weights_proj", 1),
-            ]
-            stacked_params_mapping.extend(indexer_fused_mapping)
+        # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+        indexer_fused_mapping = [
+            ("wk_weights_proj", "wk", 0),
+            ("wk_weights_proj", "weights_proj", 1),
+        ]
+        stacked_params_mapping.extend(indexer_fused_mapping)
 
         expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
             self,
@@ -271,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        _pending_wk_fp8: dict = {}  # FP8 indexer wk dequant buffer
         for name, loaded_weight in weights:
             if "rotary_emb.inv_freq" in name:
                 continue
@@ -281,6 +278,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
                 rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
             )
             name = self._rewrite_spec_layer_name(spec_layer, name)
+
+            if _try_load_fp8_indexer_wk(
+                name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
+            ):
+                continue
+
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 # Skip non-stacked layers and experts (experts handled below).
                 if weight_name not in name:

vllm/model_executor/models/deepseek_v2.py

@@ -66,6 +66,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
 )
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    scaled_dequantize,
+)
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.sparse_attn_indexer import (
     SparseAttnIndexer,
@@ -628,10 +632,6 @@ class Indexer(nn.Module):
         self.vllm_config = vllm_config
         self.config = config
         self.quant_config = quant_config
-        self.is_fp4_ckpt = (
-            self.quant_config is not None
-            and self.quant_config.get_name() == "modelopt_fp4"
-        )
         # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
         self.topk_tokens = config.index_topk
         self.n_head = config.index_n_heads  # 64
@@ -646,36 +646,16 @@ class Indexer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.wq_b",
         )
-        if self.is_fp4_ckpt:
-            # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
-            # weights_proj does not get quantized,
-            # so we run both with quant_config=None
-            # wk may be upcasted from the default quant;
-            # experiments show fusion is always faster unless WK proj is in FP4,
-            # which is not the case for all known quants.
-            self.wk_weights_proj = MergedColumnParallelLinear(
-                hidden_size,
-                [self.head_dim, self.n_head],
-                bias=False,
-                quant_config=None,
-                disable_tp=True,
-                prefix=f"{prefix}.wk_weights_proj",
-            )
-        else:
-            self.wk = ReplicatedLinear(
-                hidden_size,
-                self.head_dim,
-                bias=False,
-                quant_config=quant_config,
-                prefix=f"{prefix}.wk",
-            )
-            self.weights_proj = ReplicatedLinear(
-                hidden_size,
-                self.n_head,
-                bias=False,
-                quant_config=None,
-                prefix=f"{prefix}.weights_proj",
-            )
+        # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
+        # FP8 wk weights are upcasted to BF16 during loading to maintain fusion.
+        self.wk_weights_proj = MergedColumnParallelLinear(
+            hidden_size,
+            [self.head_dim, self.n_head],
+            bias=False,
+            quant_config=None,
+            disable_tp=True,
+            prefix=f"{prefix}.wk_weights_proj",
+        )
         self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
         self.softmax_scale = self.head_dim**-0.5
 
@@ -716,14 +696,10 @@ class Indexer(nn.Module):
         q_pe, q_nope = torch.split(
             q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
         )
-        if self.is_fp4_ckpt:
-            # Fused wk + weights_proj: one GEMM, then split
-            kw, _ = self.wk_weights_proj(hidden_states)
-            k = kw[:, : self.head_dim]
-            weights = kw[:, self.head_dim :]
-        else:
-            k, _ = self.wk(hidden_states)
-            weights, _ = self.weights_proj(hidden_states)
+        # Fused wk + weights_proj: one GEMM, then split
+        kw, _ = self.wk_weights_proj(hidden_states)
+        k = kw[:, : self.head_dim]
+        weights = kw[:, self.head_dim :]
 
         k = self.k_norm(k)
         k_pe, k_nope = torch.split(
@@ -761,6 +737,46 @@ class Indexer(nn.Module):
         return self.indexer_op(hidden_states, q_fp8, k, weights)
 
 
+def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
+    """
+    We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
+    in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
+    Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK
+    weights and scale, and when both are available, dequantizes to BF16 and stores into
+    the fused wk_weights_proj.weight parameter.
+    """
+    if "indexer.wk." not in name or "wk_weights" in name:
+        return False  # Weight is not an isolated WK weight for the indexer, ignore.
+    is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
+    is_scale = "weight_scale_inv" in name
+    if not is_weight and not is_scale:
+        return False  # WK is not in FP8 format, ignore.
+    # Buffer this tensor (weight or scale) until both have arrived.
+    layer_prefix = name.rsplit(".wk.", 1)[0]  # e.g. "model.layers.0.self_attn.indexer"
+    entry = buf.setdefault(layer_prefix, {})
+    entry["weight" if is_weight else "scale"] = tensor
+    if "weight" not in entry or "scale" not in entry:
+        return True  # still waiting for the other param
+
+    # We have both weight and scale: dequantize FP8 to BF16.
+    weight_fp8, scale_inv = entry["weight"], entry["scale"]
+    del buf[layer_prefix]
+    block_size = weight_fp8.shape[1] // scale_inv.shape[1]
+    weight_bf16 = scaled_dequantize(
+        weight_fp8,
+        scale_inv,
+        group_shape=GroupShape(block_size, block_size),
+        out_dtype=torch.bfloat16,
+    )
+
+    # Load the dequantized weight into shard 0 of the fused buffer.
+    fused_name = f"{layer_prefix}.wk_weights_proj.weight"
+    param = params_dict[fused_name]
+    param.weight_loader(param, weight_bf16, 0)
+    loaded_params.add(fused_name)
+    return True
+
+
 def _min_latency_fused_qkv_a_proj_impl(
     input_: torch.Tensor,
     weight: torch.Tensor,
@@ -1344,10 +1360,6 @@ class DeepseekV2ForCausalLM(
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
-        self.is_fp4_ckpt = (
-            self.quant_config is not None
-            and self.quant_config.get_name() == "modelopt_fp4"
-        )
 
         qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
         qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
@@ -1473,13 +1485,13 @@ class DeepseekV2ForCausalLM(
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
         ]
-        if self.is_fp4_ckpt:
-            # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
-            indexer_fused_mapping = [
-                ("wk_weights_proj", "wk", 0),
-                ("wk_weights_proj", "weights_proj", 1),
-            ]
-            stacked_params_mapping.extend(indexer_fused_mapping)
+        # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
+        _pending_wk_fp8: dict = {}  # When WK is in FP8, we dequant to BF16 for fusion
+        indexer_fused_mapping = [
+            ("wk_weights_proj", "wk", 0),
+            ("wk_weights_proj", "weights_proj", 1),
+        ]
+        stacked_params_mapping.extend(indexer_fused_mapping)
 
         if self.use_mha:
             stacked_params_mapping.extend(mha_params_mapping)
@@ -1516,6 +1528,11 @@ class DeepseekV2ForCausalLM(
                 rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
             )
 
+            if _try_load_fp8_indexer_wk(
+                name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
+            ):
+                continue
+
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 # Skip non-stacked layers and experts (experts handled below).
                 if weight_name not in name:

vllm/model_executor/models/gemma4_mm.py

@@ -1254,9 +1254,10 @@ class Gemma4ForConditionalGeneration(
             # computation (using token_type_ids == 0 as text_mask).
             # Replicate this: map image token positions to token 0.
             if is_multimodal is not None:
-                is_multimodal = is_multimodal.to(input_ids.device)
                 ple_input_ids = torch.where(
-                    is_multimodal, torch.zeros_like(input_ids), input_ids
+                    is_multimodal.to(input_ids.device, non_blocking=True),
+                    torch.zeros_like(input_ids),
+                    input_ids,
                 )
             else:
                 ple_input_ids = input_ids

vllm/model_executor/models/registry.py

@@ -1302,15 +1302,6 @@ ModelRegistry = _ModelRegistry(
     }
 )
 
-if envs.VLLM_USE_SPECIALIZED_MODELS:
-    from vllm.model_executor.specialized_models import get_specialized_models
-
-    for _arch, (_mod, _cls) in get_specialized_models().items():
-        ModelRegistry.models[_arch] = _LazyRegisteredModel(
-            module_name=_mod,
-            class_name=_cls,
-        )
-
 _T = TypeVar("_T")
 
 

vllm/model_executor/specialized_models/README.md

@@ -1,19 +0,0 @@
-# [Experimental] Specialized Models
-
-This directory contains experimental, hand-tuned implementations for a small number of selected models. Each subdirectory targets a specific combination of model architecture (including all tensor shapes), quantization scheme, attention backend, and hardware.
-
-For example, `deepseek_v3_2_nvfp4/` targets `nvidia/DeepSeek-V3.2-NVFP4` with FP8 FlashInfer sparse MLA on Blackwell GPUs.
-
-**To opt in, set `VLLM_USE_SPECIALIZED_MODELS=1`.** When enabled, vLLM will prefer a specialized implementation over the generic one if a match is available.
-
-## Development Philosophy
-
-These implementations prioritize iteration speed and checkpoint-specific performance over broad reuse. They may target a very narrow use case and are not expected to cover the full vLLM feature surface. Known limitations include:
-
-- Parallelism strategy support may be incomplete (e.g. TP only, no EP, or vice versa).
-- `torch.compile` compatibility may be limited or untested.
-- Behavior with checkpoint formats outside the intended target is unsupported.
-
-Also, code duplication across implementations is intentional — each model should be free to evolve and be optimized independently without risk of regressing another.
-
-Code here is experimental and may be short-lived. Generic features and anything intended for long-term support should live in `../models/`.

vllm/model_executor/specialized_models/__init__.py

@@ -1,36 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Specialized model implementations.
-
-Each entry maps a vLLM architecture name to a (module_path, class_name)
-tuple, exactly like the main model registry.  When
-``VLLM_USE_SPECIALIZED_MODELS=1`` the main registry merges these entries
-so they take priority over the generic implementations.
-
-To add a new specialized model:
-  1. Create a sub-package under this directory.
-  2. Add the architecture -> (module, class) mapping to ``_MODELS`` below.
-"""
-
-from __future__ import annotations
-
-# ── Model list ───────────────────────────────────────────────────────
-# Maps architecture name -> (fully-qualified module, class name).
-# When the flag is enabled, these override the corresponding entries
-# in the main registry.
-_MODELS: dict[str, tuple[str, str]] = {
-    "DeepseekV32ForCausalLM": (
-        "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4",
-        "DeepseekV32ForCausalLM",
-    ),
-    "DeepSeekMTPModel": (
-        "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4",
-        "DeepSeekMTP",
-    ),
-}
-
-
-def get_specialized_models() -> dict[str, tuple[str, str]]:
-    """Return the specialized model registry."""
-    return _MODELS

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/README.md

@@ -1,34 +0,0 @@
-# nvidia/DeepSeek-V3.2-NVFP4
-
-An optimized implementation for `nvidia/DeepSeek-V3.2-NVFP4` with FP8 FlashInfer MLA on Blackwell GPUs.
-
-The main win comes from aggressively fusing ops in the attention path, across the MLA and sparse-indexer boundary, which is critical for low latency.
-On top of manual fusions, the implementation uses `torch.compile` with vLLM's custom fusion passes to fuse remaining miscellaneous ops.
-It is compatible with piecewise CUDA graphs for prefill and full CUDA graphs for decode.
-
-TP and EP are supported; PP is not.
-MTP is supported.
-
-## Usage
-
-```bash
-export VLLM_USE_SPECIALIZED_MODELS=1
-export VLLM_USE_V2_MODEL_RUNNER=1
-export TRTLLM_ENABLE_PDL=1
-
-NUM_GPUS=4
-
-# With TP
-vllm serve nvidia/DeepSeek-V3.2-NVFP4 \
-    -tp 4 \
-    --compilation-config '{"max_cudagraph_capture_size": 1024}' \
-    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
-    --kernel-config.enable_flashinfer_autotune=False
-
-# With attention DP + MoE EP
-vllm serve nvidia/DeepSeek-V3.2-NVFP4 \
-    -dp $NUM_GPUS -ep \
-    --compilation-config '{"max_cudagraph_capture_size": 1024}' \
-    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
-    --kernel-config.enable_flashinfer_autotune=False
-```

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/__init__.py

@@ -1,8 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""DeepSeek V3.2 model optimized for SM100 (Blackwell)."""
-
-from .model import DeepseekV32ForCausalLM
-from .mtp import DeepSeekMTP
-
-__all__ = ["DeepseekV32ForCausalLM", "DeepSeekMTP"]

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/kernels.py

@@ -1,931 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import torch
-
-from vllm.platforms import current_platform
-from vllm.triton_utils import tl, triton
-from vllm.utils.torch_utils import direct_register_custom_op
-
-
-@triton.jit
-def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
-    x = x.to(tl.float32)
-    mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
-    rrms = tl.rsqrt(mean_sq + eps)
-    w = w.to(tl.float32)
-    return (x * rrms) * w
-
-
-@triton.jit
-def _fused_mtp_entry_kernel(
-    inputs_embeds_ptr,
-    inputs_embeds_stride,
-    hidden_states_ptr,
-    hidden_states_stride,
-    positions_ptr,
-    enorm_weight_ptr,
-    hnorm_weight_ptr,
-    out_ptr,
-    out_stride,
-    e_eps,
-    h_eps,
-    HIDDEN_SIZE: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    tok_idx = tl.program_id(0)
-    which = tl.program_id(1)  # 0: enorm, 1: hnorm
-
-    offs = tl.arange(0, BLOCK_SIZE)
-    mask = offs < HIDDEN_SIZE
-
-    if which == 0:
-        position = tl.load(positions_ptr + tok_idx)
-        x = tl.load(
-            inputs_embeds_ptr + tok_idx * inputs_embeds_stride + offs,
-            mask=mask,
-            other=0.0,
-        ).to(tl.float32)
-        # Mask out inputs_embeds when position == 0 (MTP convention).
-        keep = (position != 0).to(tl.float32)
-        x = x * keep
-        w = tl.load(enorm_weight_ptr + offs, mask=mask).to(tl.float32)
-        mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
-        rrms = tl.rsqrt(mean_sq + e_eps)
-        y = (x * rrms) * w
-        tl.store(
-            out_ptr + tok_idx * out_stride + offs,
-            y,
-            mask=mask,
-        )
-    else:
-        h = tl.load(
-            hidden_states_ptr + tok_idx * hidden_states_stride + offs,
-            mask=mask,
-            other=0.0,
-        ).to(tl.float32)
-        w = tl.load(hnorm_weight_ptr + offs, mask=mask).to(tl.float32)
-        mean_sq = tl.sum(h * h, axis=0) / HIDDEN_SIZE
-        rrms = tl.rsqrt(mean_sq + h_eps)
-        y = (h * rrms) * w
-        tl.store(
-            out_ptr + tok_idx * out_stride + HIDDEN_SIZE + offs,
-            y,
-            mask=mask,
-        )
-
-
-@triton.jit
-def _fused_mtp_entry_eps_kernel(
-    inputs_embeds_ptr,
-    inputs_embeds_stride,
-    hidden_states_ptr,
-    hidden_states_stride,
-    positions_ptr,
-    enorm_weight_ptr,
-    hnorm_weight_ptr,
-    e_eps_ptr,
-    h_eps_ptr,
-    out_ptr,
-    out_stride,
-    HIDDEN_SIZE: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """Same as _fused_mtp_entry_kernel but reads eps from 0-dim tensors."""
-    tok_idx = tl.program_id(0)
-    which = tl.program_id(1)
-
-    offs = tl.arange(0, BLOCK_SIZE)
-    mask = offs < HIDDEN_SIZE
-
-    if which == 0:
-        position = tl.load(positions_ptr + tok_idx)
-        x = tl.load(
-            inputs_embeds_ptr + tok_idx * inputs_embeds_stride + offs,
-            mask=mask,
-            other=0.0,
-        ).to(tl.float32)
-        keep = (position != 0).to(tl.float32)
-        x = x * keep
-        w = tl.load(enorm_weight_ptr + offs, mask=mask).to(tl.float32)
-        mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
-        e_eps = tl.load(e_eps_ptr)
-        rrms = tl.rsqrt(mean_sq + e_eps)
-        y = (x * rrms) * w
-        tl.store(
-            out_ptr + tok_idx * out_stride + offs,
-            y,
-            mask=mask,
-        )
-    else:
-        h = tl.load(
-            hidden_states_ptr + tok_idx * hidden_states_stride + offs,
-            mask=mask,
-            other=0.0,
-        ).to(tl.float32)
-        w = tl.load(hnorm_weight_ptr + offs, mask=mask).to(tl.float32)
-        mean_sq = tl.sum(h * h, axis=0) / HIDDEN_SIZE
-        h_eps = tl.load(h_eps_ptr)
-        rrms = tl.rsqrt(mean_sq + h_eps)
-        y = (h * rrms) * w
-        tl.store(
-            out_ptr + tok_idx * out_stride + HIDDEN_SIZE + offs,
-            y,
-            mask=mask,
-        )
-
-
-def _fused_mtp_entry_impl(
-    inputs_embeds: torch.Tensor,
-    hidden_states: torch.Tensor,
-    positions: torch.Tensor,
-    enorm_weight: torch.Tensor,
-    hnorm_weight: torch.Tensor,
-    e_eps: torch.Tensor,
-    h_eps: torch.Tensor,
-    out: torch.Tensor,
-) -> torch.Tensor:
-    num_tokens, hidden_size = inputs_embeds.shape
-    BLOCK_SIZE = triton.next_power_of_2(hidden_size)
-    _fused_mtp_entry_eps_kernel[(num_tokens, 2)](
-        inputs_embeds,
-        inputs_embeds.stride(0),
-        hidden_states,
-        hidden_states.stride(0),
-        positions,
-        enorm_weight,
-        hnorm_weight,
-        e_eps,
-        h_eps,
-        out,
-        out.stride(0),
-        HIDDEN_SIZE=hidden_size,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=8,
-    )
-    return out
-
-
-def _fused_mtp_entry_fake(
-    inputs_embeds: torch.Tensor,
-    hidden_states: torch.Tensor,
-    positions: torch.Tensor,
-    enorm_weight: torch.Tensor,
-    hnorm_weight: torch.Tensor,
-    e_eps: torch.Tensor,
-    h_eps: torch.Tensor,
-    out: torch.Tensor,
-) -> torch.Tensor:
-    del (
-        inputs_embeds,
-        hidden_states,
-        positions,
-        enorm_weight,
-        hnorm_weight,
-        e_eps,
-        h_eps,
-    )
-    return out
-
-
-direct_register_custom_op(
-    op_name="fused_mtp_entry",
-    op_func=_fused_mtp_entry_impl,
-    fake_impl=_fused_mtp_entry_fake,
-    mutates_args=["out"],
-    dispatch_key=current_platform.dispatch_key,
-)
-
-
-def fused_mtp_entry(
-    inputs_embeds: torch.Tensor,
-    hidden_states: torch.Tensor,
-    positions: torch.Tensor,
-    enorm_weight: torch.Tensor,
-    hnorm_weight: torch.Tensor,
-    e_eps: torch.Tensor,
-    h_eps: torch.Tensor,
-) -> torch.Tensor:
-    """Fused: mask(pos==0) + enorm(embeds) | hnorm(hidden) -> concat.
-
-    Output is the concatenation [enorm(embeds), hnorm(hidden)] in the
-    last dim, ready to feed into eh_proj. `e_eps`/`h_eps` are 0-dim fp32
-    tensors (not Python floats) so the custom op stays tensor-only.
-    """
-    num_tokens, hidden_size = inputs_embeds.shape
-    out = torch.empty(
-        num_tokens,
-        hidden_size * 2,
-        dtype=inputs_embeds.dtype,
-        device=inputs_embeds.device,
-    )
-    return torch.ops.vllm.fused_mtp_entry(
-        inputs_embeds,
-        hidden_states,
-        positions,
-        enorm_weight,
-        hnorm_weight,
-        e_eps,
-        h_eps,
-        out,
-    )
-
-
-@triton.jit
-def _layer_norm(x, w, b, eps, mask, HIDDEN_SIZE: tl.constexpr):
-    x = x.to(tl.float32)
-    mean = tl.sum(x, axis=0) / HIDDEN_SIZE
-    diff = tl.where(mask, x - mean, 0.0)
-    var = tl.sum(diff * diff, axis=0) / HIDDEN_SIZE
-    rstd = tl.rsqrt(var + eps)
-
-    w = w.to(tl.float32)
-    b = b.to(tl.float32)
-    return (x - mean) * rstd * w + b
-
-
-@triton.jit
-def _rope(
-    base_ptr,
-    head_stride,
-    cos,
-    sin,
-    NUM_HEADS: tl.constexpr,
-    HALF_ROT_DIM: tl.constexpr,
-    START_OFFSET: tl.constexpr,
-    INTERLEAVED: tl.constexpr,
-):
-    head_offset = tl.arange(0, NUM_HEADS)
-    dim_offset = tl.arange(0, HALF_ROT_DIM)
-    base_ptr = base_ptr + head_offset[:, None] * head_stride + START_OFFSET
-    if INTERLEAVED:
-        x1 = tl.load(base_ptr + dim_offset * 2).to(tl.float32)
-        x2 = tl.load(base_ptr + dim_offset * 2 + 1).to(tl.float32)
-        tl.store(base_ptr + dim_offset * 2, x1 * cos - x2 * sin)
-        tl.store(base_ptr + dim_offset * 2 + 1, x2 * cos + x1 * sin)
-    else:
-        x1 = tl.load(base_ptr + dim_offset).to(tl.float32)
-        x2 = tl.load(base_ptr + dim_offset + HALF_ROT_DIM).to(tl.float32)
-        tl.store(base_ptr + dim_offset, x1 * cos - x2 * sin)
-        tl.store(base_ptr + dim_offset + HALF_ROT_DIM, x2 * cos + x1 * sin)
-
-
-@triton.jit
-def _get_cos_sin(
-    cos_sin_cache_ptr,
-    cos_sin_cache_stride,
-    pos,
-    HALF_ROT_DIM: tl.constexpr,
-):
-    block = tl.arange(0, HALF_ROT_DIM)
-    cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block)
-    cos = cos.to(tl.float32)
-    sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM)
-    sin = sin.to(tl.float32)
-    return cos, sin
-
-
-@triton.jit
-def _fp8_ue8m0_quantize(vals):
-    """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
-
-    Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
-    """
-    vals = vals.to(tl.float32)
-    amax = tl.max(tl.abs(vals))
-    scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
-    scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
-    fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
-    return fp8_vals, scale
-
-
-@triton.jit
-def _fp8_quant_and_cache_write(
-    vals,
-    mask,
-    slot_idx,
-    kv_cache_ptr,
-    kv_cache_scale_ptr,
-    cache_block_size,
-    cache_stride,
-    offsets,
-    HEAD_DIM: tl.constexpr,
-):
-    k_fp8, scale = _fp8_ue8m0_quantize(vals)
-
-    block_idx = slot_idx // cache_block_size
-    block_offset = slot_idx % cache_block_size
-    block_start = block_idx * cache_block_size * cache_stride
-
-    tl.store(
-        kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
-        k_fp8,
-        mask=mask,
-    )
-    scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
-    tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
-
-
-@triton.jit
-def _fused_norm_rope_kernel(
-    pos_ptr,
-    # Q RMS norm
-    q_c_ptr,
-    q_c_stride,
-    q_rms_norm_w_ptr,
-    q_rms_eps,
-    q_c_out_ptr,
-    q_c_out_stride,
-    Q_DIM: tl.constexpr,
-    Q_BLOCK_SIZE: tl.constexpr,
-    # KV RMS norm
-    kv_ptr,
-    kv_stride,
-    kv_rms_norm_w_ptr,
-    kv_rms_eps,
-    KV_DIM: tl.constexpr,
-    # KV RoPE
-    kpe_ptr,
-    kpe_stride,
-    kpe_rope_cos_sin_cache_ptr,
-    kpe_rope_cos_sin_cache_stride,
-    KPE_HALF_ROT_DIM: tl.constexpr,
-    # Index K layer norm
-    index_k_ptr,
-    index_k_stride,
-    index_k_layer_norm_w_ptr,
-    index_k_layer_norm_bias_ptr,
-    index_k_layer_norm_eps,
-    INDEX_K_DIM: tl.constexpr,
-    INDEX_K_BLOCK_SIZE: tl.constexpr,
-    # Index K RoPE
-    index_k_rope_cos_sin_cache_ptr,
-    index_k_rope_cos_sin_cache_stride,
-    INDEX_K_HALF_ROT_DIM: tl.constexpr,
-    # Index K fp32 scratch buffer for layernorm → RoPE handoff
-    index_k_normed_ptr,
-    # Cache params (shared by indexer K and MLA)
-    slot_mapping_ptr,
-    # Index K FP8 cache
-    indexer_cache_ptr,
-    indexer_cache_scale_ptr,
-    indexer_cache_block_size,
-    indexer_cache_stride,
-    # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
-    mla_cache_ptr,
-    mla_cache_block_stride,
-    mla_cache_entry_stride,
-    MLA_CACHE_FP8: tl.constexpr,
-    mla_cache_scale_ptr,
-    # Top k indices
-    topk_indices_ptr,
-    topk_indices_stride,
-    TOPK: tl.constexpr,
-    TOPK_BLOCK_SIZE: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    tok_idx = tl.program_id(1)
-    if pid == 3:
-        # Fill top k indices buffer with -1
-        for i in range(0, TOPK, TOPK_BLOCK_SIZE):
-            offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
-            mask = offset < TOPK
-            tl.store(
-                topk_indices_ptr + tok_idx * topk_indices_stride + offset,
-                -1,
-                mask=mask,
-            )
-        return
-
-    if slot_mapping_ptr is None:
-        # Memory profiling run.
-        return
-    slot_idx = tl.load(slot_mapping_ptr + tok_idx)
-    if slot_idx < 0:
-        # Padding
-        return
-
-    if pid == 2:
-        # Q RMS norm
-        q_block = tl.arange(0, Q_BLOCK_SIZE)
-        q_mask = q_block < Q_DIM
-        q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
-        q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
-        q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
-        tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
-    elif pid == 1:
-        # KV RMS Norm + KV RoPE + MLA concat_and_cache.
-        # Merged so the normed kv_c and RoPE'd k_pe can be written
-        # to the MLA KV cache directly without a separate kernel.
-
-        # KV RMS Norm (result stays in registers for MLA cache write)
-        kv_block = tl.arange(0, KV_DIM)
-        kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block)
-        kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block)
-        kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM)
-
-        # KV RoPE (interleaved) on k_pe — in registers only.
-        # k_pe is not needed after the cache write (MLA decode reads
-        # from kv_cache), so we skip writing back to kpe_ptr.
-        pos = tl.load(pos_ptr + tok_idx)
-        cos, sin = _get_cos_sin(
-            kpe_rope_cos_sin_cache_ptr,
-            kpe_rope_cos_sin_cache_stride,
-            pos,
-            KPE_HALF_ROT_DIM,
-        )
-        dim_off = tl.arange(0, KPE_HALF_ROT_DIM)
-        kpe_base = kpe_ptr + tok_idx * kpe_stride
-        x1 = tl.load(kpe_base + dim_off * 2).to(tl.float32)
-        x2 = tl.load(kpe_base + dim_off * 2 + 1).to(tl.float32)
-        r1 = x1 * cos - x2 * sin
-        r2 = x2 * cos + x1 * sin
-
-        # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.
-        if mla_cache_entry_stride == 0:
-            return
-
-        mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
-        mla_block_idx = slot_idx // mla_block_size
-        mla_block_off = slot_idx % mla_block_size
-        dst = (
-            mla_cache_ptr
-            + mla_block_idx * mla_cache_block_stride
-            + mla_block_off * mla_cache_entry_stride
-        )
-        # kv_c_normed (KV_DIM elements)
-        if MLA_CACHE_FP8:
-            scale = tl.load(mla_cache_scale_ptr)
-            kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)
-            tl.store(dst + kv_block, kv_c_fp8)
-        else:
-            tl.store(dst + kv_block, kv_c)
-        # k_pe_roped (from registers, interleaved layout)
-        if MLA_CACHE_FP8:
-            tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))
-            tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv))
-        else:
-            tl.store(dst + KV_DIM + dim_off * 2, r1)
-            tl.store(dst + KV_DIM + dim_off * 2 + 1, r2)
-    elif pid == 0:
-        # Fused: Index K LayerNorm + RoPE + FP8 quant + cache write.
-        # Eliminates the separate indexer_k_quant_and_cache kernel launch.
-
-        # 1. LayerNorm → fp32 temp buffer
-        index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE)
-        index_k_mask = index_k_block < INDEX_K_DIM
-        index_k = tl.load(
-            index_k_ptr + tok_idx * index_k_stride + index_k_block,
-            mask=index_k_mask,
-            other=0.0,
-        )
-        index_k_w = tl.load(index_k_layer_norm_w_ptr + index_k_block, mask=index_k_mask)
-        index_k_b = tl.load(
-            index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask
-        )
-        normed = _layer_norm(
-            index_k,
-            index_k_w,
-            index_k_b,
-            index_k_layer_norm_eps,
-            index_k_mask,
-            INDEX_K_DIM,
-        )
-        # Write to a fp32 scratch buffer so RoPE can read the two
-        # halves without Triton pointer-aliasing issues.
-        scratch = index_k_normed_ptr + tok_idx * INDEX_K_DIM
-        tl.store(scratch + index_k_block, normed, mask=index_k_mask)
-
-        # 2. RoPE (neox / non-interleaved) on the full vector.
-        pos = tl.load(pos_ptr + tok_idx)
-        cos_full = tl.load(
-            index_k_rope_cos_sin_cache_ptr
-            + pos * index_k_rope_cos_sin_cache_stride
-            + index_k_block % INDEX_K_HALF_ROT_DIM,
-            mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
-            other=1.0,
-        ).to(tl.float32)
-        sin_full = tl.load(
-            index_k_rope_cos_sin_cache_ptr
-            + pos * index_k_rope_cos_sin_cache_stride
-            + INDEX_K_HALF_ROT_DIM
-            + index_k_block % INDEX_K_HALF_ROT_DIM,
-            mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
-            other=0.0,
-        ).to(tl.float32)
-        # XOR with HALF swaps the first/second half of the rotation
-        # region to get each element's partner.
-        partner_offs = tl.where(
-            index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
-            index_k_block ^ INDEX_K_HALF_ROT_DIM,
-            index_k_block,
-        )
-        full = tl.load(scratch + index_k_block, mask=index_k_mask)
-        # Atomic read for the partner: tl.atomic_add(ptr, 0) returns the
-        # current value with guaranteed store visibility, avoiding the
-        # Triton compiler's aliasing issue with different offset expressions.
-        zeros = tl.zeros([INDEX_K_BLOCK_SIZE], dtype=tl.float32)
-        partner = tl.atomic_add(scratch + partner_offs, zeros, mask=index_k_mask)
-        sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0)
-        roped = full * cos_full + sign * partner * sin_full
-        result = tl.where(index_k_block < 2 * INDEX_K_HALF_ROT_DIM, roped, full)
-
-        # 3. FP8 quantize + cache write from registers.
-        #    No need to write back to index_k_ptr — the only consumer
-        #    (sparse_attn_indexer) reads from the cache, not index_k.
-        _fp8_quant_and_cache_write(
-            result,
-            index_k_mask,
-            slot_idx,
-            indexer_cache_ptr,
-            indexer_cache_scale_ptr,
-            indexer_cache_block_size,
-            indexer_cache_stride,
-            index_k_block,
-            INDEX_K_DIM,
-        )
-
-
-def fused_norm_rope(
-    positions: torch.Tensor,
-    q_c: torch.Tensor,
-    q_rms_norm_w: torch.Tensor,
-    q_rms_eps: float,
-    kv_c: torch.Tensor,
-    kv_rms_norm_w: torch.Tensor,
-    kv_rms_eps: float,
-    k_pe: torch.Tensor,
-    k_rope_cos_sin_cache: torch.Tensor,
-    index_k: torch.Tensor,
-    index_k_layer_norm_w: torch.Tensor,
-    index_k_layer_norm_bias: torch.Tensor,
-    index_k_layer_norm_eps: float,
-    index_k_rope_cos_sin_cache: torch.Tensor,
-    topk_indices_buffer: torch.Tensor,
-    # Cache params for fused writes (single slot_mapping for both caches)
-    slot_mapping: torch.Tensor | None = None,
-    indexer_k_cache: torch.Tensor | None = None,
-    mla_kv_cache: torch.Tensor | None = None,
-    mla_kv_cache_dtype: str = "auto",
-    mla_k_scale: torch.Tensor | None = None,
-) -> torch.Tensor:
-    assert positions.ndim == 1
-    assert q_c.ndim == 2
-    assert kv_c.ndim == 2
-    assert k_pe.ndim == 2
-    assert index_k.ndim == 2
-    assert topk_indices_buffer.ndim == 2
-
-    num_tokens = positions.shape[0]
-    q_dim = q_c.shape[-1]
-    kv_dim = kv_c.shape[-1]
-    index_k_dim = index_k.shape[-1]
-    topk = topk_indices_buffer.shape[-1]
-    device = positions.device
-
-    # --- Indexer K cache setup ---
-    if indexer_k_cache is not None:
-        assert slot_mapping is not None
-        idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
-        idx_cache_block_size = indexer_k_cache.shape[1]
-        idx_cache_stride = indexer_k_cache.shape[2]
-        if indexer_k_cache.dtype == torch.uint8:
-            indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
-    else:
-        idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device)
-        indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device)
-        slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
-        idx_cache_block_size = 1
-        idx_cache_stride = 1
-
-    # --- MLA KV cache setup ---
-    mla_cache_fp8 = mla_kv_cache_dtype != "auto"
-    if mla_kv_cache is not None:
-        mla_block_stride = mla_kv_cache.stride(0)
-        mla_entry_stride = mla_kv_cache.stride(1)
-        if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
-            mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
-        if mla_k_scale is None:
-            mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
-    else:
-        # Dummy values — pid 2 will skip the MLA cache write because
-        # slot_mapping is all -1.
-        mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
-        mla_block_stride = 0
-        mla_entry_stride = 0
-        mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
-
-    # fp32 scratch buffer for layernorm output → RoPE handoff.
-    index_k_normed = torch.empty(
-        num_tokens, index_k_dim, dtype=torch.float32, device=device
-    )
-
-    q_c_out = torch.empty_like(q_c)
-    _fused_norm_rope_kernel[(4, num_tokens)](
-        positions,
-        # Q RMS norm
-        q_c,
-        q_c.stride(0),
-        q_rms_norm_w,
-        q_rms_eps,
-        q_c_out,
-        q_c_out.stride(0),
-        q_dim,
-        triton.next_power_of_2(q_dim),
-        # KV RMS norm
-        kv_c,
-        kv_c.stride(0),
-        kv_rms_norm_w,
-        kv_rms_eps,
-        kv_dim,
-        # KV RoPE
-        k_pe,
-        k_pe.stride(0),
-        k_rope_cos_sin_cache,
-        k_rope_cos_sin_cache.stride(0),
-        k_rope_cos_sin_cache.shape[-1] // 2,
-        # Index K layer norm + RoPE + FP8 quant
-        index_k,
-        index_k.stride(0),
-        index_k_layer_norm_w,
-        index_k_layer_norm_bias,
-        index_k_layer_norm_eps,
-        index_k_dim,
-        triton.next_power_of_2(index_k_dim),
-        index_k_rope_cos_sin_cache,
-        index_k_rope_cos_sin_cache.stride(0),
-        index_k_rope_cos_sin_cache.shape[-1] // 2,
-        index_k_normed,
-        # Cache params
-        slot_mapping,
-        indexer_k_cache,
-        idx_cache_scale_view,
-        idx_cache_block_size,
-        idx_cache_stride,
-        # MLA KV cache (uses same slot_mapping)
-        mla_kv_cache,
-        mla_block_stride,
-        mla_entry_stride,
-        mla_cache_fp8,
-        mla_k_scale,
-        # Top k indices buffer
-        topk_indices_buffer,
-        topk_indices_buffer.stride(0),
-        topk,
-        TOPK_BLOCK_SIZE=1024,
-    )
-    return q_c_out
-
-
-@triton.jit
-def _fused_q_kernel(
-    pos_ptr,
-    # MQA query PE: RoPE + FP8 pack into output tail
-    q_pe_ptr,
-    q_pe_stride0,
-    q_pe_stride1,
-    NUM_Q_HEADS: tl.constexpr,
-    q_pe_cos_sin_ptr,
-    q_pe_cos_sin_stride,
-    Q_PE_HALF_ROT_DIM: tl.constexpr,
-    # Index Q RoPE
-    index_q_ptr,
-    index_q_stride0,
-    index_q_stride1,
-    NUM_INDEX_Q_HEADS: tl.constexpr,
-    index_q_cos_sin_ptr,
-    index_q_cos_sin_stride,
-    INDEX_Q_HALF_ROT_DIM: tl.constexpr,
-    # Index Q Quantize
-    index_q_fp8_ptr,
-    index_q_fp8_stride0,
-    index_q_fp8_stride1,
-    INDEX_Q_HEAD_DIM: tl.constexpr,
-    # MQA query pack: quantize ql_nope and RoPE+quantize q_pe into mqa_q_fp8
-    ql_nope_ptr,
-    ql_nope_stride0,
-    ql_nope_stride1,
-    mqa_q_fp8_ptr,
-    mqa_q_fp8_stride0,
-    mqa_q_fp8_stride1,
-    q_scale_ptr,
-    QL_NOPE_DIM: tl.constexpr,
-    QL_NOPE_BLOCK: tl.constexpr,
-    # Index weights
-    index_weights_ptr,
-    index_weights_stride,
-    index_weights_softmax_scale,
-    index_weights_head_scale,
-    index_weights_out_ptr,
-    index_weights_out_stride,
-):
-    pid = tl.program_id(0)
-    tok_idx = tl.program_id(1)
-    head_idx = tl.program_id(2)
-
-    if pid == 2:
-        # ql_nope quantize + pack into the front of mqa_q_fp8.
-        if 2 * head_idx >= NUM_Q_HEADS:
-            return
-
-        scale = tl.load(q_scale_ptr)
-        for local_head in range(2):
-            q_head_idx = head_idx * 2 + local_head
-            if q_head_idx < NUM_Q_HEADS:
-                ql_nope_off = tl.arange(0, QL_NOPE_BLOCK)
-                ql_nope_mask = ql_nope_off < QL_NOPE_DIM
-                ql_nope = tl.load(
-                    ql_nope_ptr
-                    + tok_idx * ql_nope_stride0
-                    + q_head_idx * ql_nope_stride1
-                    + ql_nope_off,
-                    mask=ql_nope_mask,
-                ).to(tl.float32)
-                ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv)
-                tl.store(
-                    mqa_q_fp8_ptr
-                    + tok_idx * mqa_q_fp8_stride0
-                    + q_head_idx * mqa_q_fp8_stride1
-                    + ql_nope_off,
-                    ql_nope_fp8,
-                    mask=ql_nope_mask,
-                )
-        return
-    elif pid == 0:
-        # q_pe RoPE + quantize + pack into the tail of mqa_q_fp8.
-        if 2 * head_idx >= NUM_Q_HEADS:
-            return
-
-        pos = tl.load(pos_ptr + tok_idx)
-        cos, sin = _get_cos_sin(
-            q_pe_cos_sin_ptr,
-            q_pe_cos_sin_stride,
-            pos,
-            Q_PE_HALF_ROT_DIM,
-        )
-
-        scale = tl.load(q_scale_ptr)
-        for local_head in range(2):
-            q_head_idx = head_idx * 2 + local_head
-            if q_head_idx < NUM_Q_HEADS:
-                rot_off = tl.arange(0, Q_PE_HALF_ROT_DIM)
-                x1 = tl.load(
-                    q_pe_ptr
-                    + tok_idx * q_pe_stride0
-                    + q_head_idx * q_pe_stride1
-                    + rot_off * 2,
-                ).to(tl.float32)
-                x2 = tl.load(
-                    q_pe_ptr
-                    + tok_idx * q_pe_stride0
-                    + q_head_idx * q_pe_stride1
-                    + rot_off * 2
-                    + 1
-                ).to(tl.float32)
-                r1 = x1 * cos - x2 * sin
-                r2 = x2 * cos + x1 * sin
-                tl.store(
-                    mqa_q_fp8_ptr
-                    + tok_idx * mqa_q_fp8_stride0
-                    + q_head_idx * mqa_q_fp8_stride1
-                    + QL_NOPE_DIM
-                    + rot_off * 2,
-                    (r1 / scale).to(tl.float8e4nv),
-                )
-                tl.store(
-                    mqa_q_fp8_ptr
-                    + tok_idx * mqa_q_fp8_stride0
-                    + q_head_idx * mqa_q_fp8_stride1
-                    + QL_NOPE_DIM
-                    + rot_off * 2
-                    + 1,
-                    (r2 / scale).to(tl.float8e4nv),
-                )
-        return
-    elif pid == 1:
-        # Index Q RoPE
-        if head_idx >= NUM_INDEX_Q_HEADS:
-            return
-
-        pos = tl.load(pos_ptr + tok_idx)
-        cos, sin = _get_cos_sin(
-            index_q_cos_sin_ptr,
-            index_q_cos_sin_stride,
-            pos,
-            INDEX_Q_HALF_ROT_DIM,
-        )
-        _rope(
-            index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1,
-            0,
-            cos,
-            sin,
-            1,
-            INDEX_Q_HALF_ROT_DIM,
-            0,
-            False,
-        )
-
-        # Index Q Quantize
-        index_q_block = tl.arange(0, INDEX_Q_HEAD_DIM)
-        index_q = tl.load(
-            index_q_ptr
-            + tok_idx * index_q_stride0
-            + head_idx * index_q_stride1
-            + index_q_block
-        )
-
-        index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
-        tl.store(
-            index_q_fp8_ptr
-            + tok_idx * index_q_fp8_stride0
-            + head_idx * index_q_fp8_stride1
-            + index_q_block,
-            index_q_fp8,
-        )
-
-        # Index weights update
-        index_weights = tl.load(
-            index_weights_ptr + tok_idx * index_weights_stride + head_idx
-        )
-        index_weights = index_weights.to(tl.float32)
-        index_weights *= index_q_scale
-        index_weights *= index_weights_softmax_scale
-        index_weights *= index_weights_head_scale
-        tl.store(
-            index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
-            index_weights,
-        )
-
-
-def fused_q(
-    positions: torch.Tensor,
-    q_pe: torch.Tensor,
-    q_pe_cos_sin_cache: torch.Tensor,
-    index_q: torch.Tensor,
-    index_q_cos_sin_cache: torch.Tensor,
-    ql_nope: torch.Tensor,
-    q_scale: torch.Tensor,
-    # Index weights
-    index_weights: torch.Tensor,
-    index_weights_softmax_scale: float,
-    index_weights_head_scale: float,
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    assert positions.ndim == 1
-    assert q_pe.ndim == 3
-    assert q_pe_cos_sin_cache.ndim == 2
-    assert index_q.ndim == 3
-    assert index_q_cos_sin_cache.ndim == 2
-
-    num_tokens = positions.shape[0]
-    num_q_heads = q_pe.shape[1]
-    num_index_q_heads = index_q.shape[1]
-    index_q_head_dim = index_q.shape[2]
-    assert ql_nope.ndim == 3
-    assert ql_nope.shape[:2] == q_pe.shape[:2]
-    mqa_q_fp8 = torch.empty(
-        q_pe.shape[0],
-        q_pe.shape[1],
-        ql_nope.shape[2] + q_pe.shape[2],
-        dtype=torch.float8_e4m3fn,
-        device=q_pe.device,
-    )
-
-    index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
-    index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
-    _fused_q_kernel[(3, num_tokens, num_index_q_heads)](
-        positions,
-        q_pe,
-        q_pe.stride(0),
-        q_pe.stride(1),
-        num_q_heads,
-        q_pe_cos_sin_cache,
-        q_pe_cos_sin_cache.stride(0),
-        q_pe_cos_sin_cache.shape[-1] // 2,
-        index_q,
-        index_q.stride(0),
-        index_q.stride(1),
-        num_index_q_heads,
-        index_q_cos_sin_cache,
-        index_q_cos_sin_cache.stride(0),
-        index_q_cos_sin_cache.shape[-1] // 2,
-        index_q_fp8,
-        index_q_fp8.stride(0),
-        index_q_fp8.stride(1),
-        index_q_head_dim,
-        ql_nope,
-        ql_nope.stride(0),
-        ql_nope.stride(1),
-        mqa_q_fp8,
-        mqa_q_fp8.stride(0),
-        mqa_q_fp8.stride(1),
-        q_scale,
-        ql_nope.shape[2],
-        triton.next_power_of_2(ql_nope.shape[2]),
-        index_weights,
-        index_weights.stride(0),
-        index_weights_softmax_scale,
-        index_weights_head_scale,
-        index_weights_out,
-        index_weights_out.stride(0),
-        num_warps=1,  # TODO: Tune this
-    )
-    return index_q_fp8, index_weights_out, mqa_q_fp8

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/layer.py

@@ -1,570 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-MLA attention and decoder layer for DeepSeek V3.2 on SM100 (Blackwell).
-
-MLAAttention:
-  KV cache update -> W_UK_T absorption -> sparse attn kernel -> W_UV up-proj
-  MLAAttention kept only as a registration stub for KV cache / backend.
-
-DecoderLayer:
-  Single decoder layer: norm -> attn -> norm -> MoE/MLP.
-"""
-
-from __future__ import annotations
-
-import torch
-from torch import nn
-
-from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
-from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.forward_context import get_forward_context
-from vllm.model_executor.layers.attention.mla_attention import MLAAttention
-from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    ReplicatedLinear,
-    RowParallelLinear,
-)
-from vllm.model_executor.layers.quantization import QuantizationConfig
-from vllm.model_executor.layers.rotary_embedding import get_rope
-from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
-from vllm.model_executor.models.deepseek_v2 import (
-    DeepseekV32IndexerCache,
-    yarn_get_mscale,
-)
-from vllm.platforms import current_platform
-from vllm.utils.torch_utils import direct_register_custom_op
-from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
-
-from .kernels import fused_norm_rope, fused_q
-from .sparse_indexer import sparse_attn_indexer
-
-
-def dsa(
-    positions: torch.Tensor,
-    q_c: torch.Tensor,
-    kv_c: torch.Tensor,
-    k_pe: torch.Tensor,
-    index_k: torch.Tensor,
-    index_weights: torch.Tensor,
-    output: torch.Tensor,
-    layer_name: str,
-) -> torch.Tensor:
-    layer = get_forward_context().no_compile_layers[layer_name]
-    attn = layer.attn
-    mla = attn.mla_attn
-
-    attn_metadata = get_forward_context().attn_metadata
-    if not isinstance(attn_metadata, dict):
-        output.zero_()
-        return output
-
-    mla_attn_metadata = attn_metadata.get(mla.layer_name)
-    if mla_attn_metadata is None:
-        output.zero_()
-        return output
-
-    num_actual_toks = mla_attn_metadata.num_actual_tokens  # type: ignore[attr-defined]
-    if num_actual_toks == 0:
-        output.zero_()
-        return output
-
-    # Step 2. fused norm + rope + cache writes
-    slot_mapping = None
-    indexer_k_cache = None
-    mla_kv_cache = None
-    mla_k_scale = None
-    idx_meta = attn_metadata.get(attn.indexer_k_cache.prefix)
-    if idx_meta is not None:
-        slot_mapping = idx_meta.slot_mapping  # type: ignore[attr-defined]
-        indexer_k_cache = attn.indexer_k_cache.kv_cache
-        mla_kv_cache = attn.mla_attn.kv_cache
-        mla_k_scale = attn.mla_attn._k_scale
-
-    q_c = fused_norm_rope(
-        positions,
-        q_c,
-        attn.q_a_layernorm_weight,
-        layer.rms_norm_eps,
-        kv_c,
-        attn.kv_a_layernorm_weight,
-        attn.rms_norm_eps,
-        k_pe,
-        attn.rotary_emb.cos_sin_cache,
-        index_k,
-        attn.indexer_k_norm.weight,
-        attn.indexer_k_norm.bias,
-        attn.rms_norm_eps,
-        attn.indexer_rope_emb.cos_sin_cache,
-        attn.topk_indices_buffer,
-        slot_mapping=slot_mapping,
-        indexer_k_cache=indexer_k_cache,
-        mla_kv_cache=mla_kv_cache,
-        mla_kv_cache_dtype=attn.mla_attn.kv_cache_dtype,
-        mla_k_scale=mla_k_scale,
-    )
-
-    # Step 3. q_c -> index_q, q
-    step3_out = torch.mm(q_c, layer._fused_step3_q_w.T)
-    index_q, q = step3_out.split(layer._q_split_sizes, dim=-1)
-    index_q = index_q.view(-1, attn.index_n_heads, attn.index_head_dim)
-    q = q.view(-1, attn.num_local_heads, attn.qk_head_dim)
-
-    # Step 4. Q RoPE + W_UK_T absorption + FP8 packing
-    q_nope, q_pe = q.split(
-        [mla.qk_nope_head_dim, mla.qk_rope_head_dim],
-        dim=-1,
-    )
-    q_nope = q_nope.transpose(0, 1)
-    ql_nope = torch.bmm(q_nope, mla.W_UK_T)
-    ql_nope = ql_nope.transpose(0, 1)
-
-    index_q_fp8, index_weights, mqa_q = fused_q(
-        positions,
-        q_pe,
-        attn.rotary_emb.cos_sin_cache,
-        index_q,
-        attn.indexer_rope_emb.cos_sin_cache,
-        ql_nope,
-        mla._q_scale,
-        index_weights,
-        attn.indexer_softmax_scale,
-        attn.index_n_heads**-0.5,
-    )
-
-    # Steps 5-6. Sparse indexer + MLA sparse decode attention
-    sparse_attn_indexer(
-        attn.indexer_k_cache.prefix,
-        attn.indexer_k_cache.kv_cache,
-        index_q_fp8,
-        index_weights,
-        attn.topk_tokens,
-        attn.index_head_dim,
-        layer.max_model_len,
-        layer.indexer_workspace_size,
-        attn.topk_indices_buffer,
-    )
-
-    mqa_q = mqa_q[:num_actual_toks]
-    kv_cache = mla.kv_cache
-    if mla.kv_cache_dtype.startswith("fp8") and mla.kv_cache_dtype != "fp8_ds_mla":
-        kv_cache = kv_cache.view(torch.float8_e4m3fn)
-    attn_out, _ = mla.impl.forward_mqa(mqa_q, kv_cache, mla_attn_metadata, mla)
-    x = attn_out.view(-1, mla.num_heads, mla.kv_lora_rank).transpose(0, 1)
-
-    out = output[:num_actual_toks].view(-1, mla.num_heads, mla.v_head_dim)
-    out = out.transpose(0, 1)
-    torch.bmm(x, mla.W_UV, out=out)
-    return output
-
-
-def dsa_fake(
-    positions: torch.Tensor,
-    q_c: torch.Tensor,
-    kv_c: torch.Tensor,
-    k_pe: torch.Tensor,
-    index_k: torch.Tensor,
-    index_weights: torch.Tensor,
-    output: torch.Tensor,
-    layer_name: str,
-) -> torch.Tensor:
-    del positions, q_c, kv_c, k_pe, index_k, index_weights, layer_name
-    return output
-
-
-direct_register_custom_op(
-    op_name="monolithic_attn",
-    op_func=dsa,
-    fake_impl=dsa_fake,
-    mutates_args=["output"],
-    dispatch_key=current_platform.dispatch_key,
-)
-
-
-class DeepseekV32DecoderLayer(nn.Module):
-    """
-    Single decoder layer: norm -> attn -> norm -> MoE/MLP.
-    Norms are raw weight + direct kernel call.
-    Gate inlined as raw weight, experts kept as FusedMoE for quantization.
-    """
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        config,
-        layer_idx: int,
-        topk_indices_buffer: torch.Tensor,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        compilation_config = get_current_vllm_config().compilation_config
-        if prefix in compilation_config.static_forward_context:
-            raise ValueError(f"Duplicate layer name: {prefix}")
-        compilation_config.static_forward_context[prefix] = self
-
-        self.layer_name = prefix
-        self.layer_idx = layer_idx
-        self.hidden_size = config.hidden_size
-        self.rms_norm_eps = config.rms_norm_eps
-        self.q_lora_rank = config.q_lora_rank
-        self.kv_lora_rank = config.kv_lora_rank
-        self.qk_rope_head_dim = config.qk_rope_head_dim
-        self.tp_size = get_tensor_model_parallel_world_size()
-
-        cache_config = vllm_config.cache_config
-        quant_config = vllm_config.quant_config
-        parallel_config = vllm_config.parallel_config
-        self.indexer_workspace_size = get_max_prefill_buffer_size(vllm_config)
-        self.max_model_len = vllm_config.model_config.max_model_len
-
-        # Use the regular vLLM RMSNorm modules so the compiler sees the
-        # canonical residual-add + RMSNorm pattern.
-        dtype = torch.get_default_dtype()
-        self.input_layernorm = RMSNorm(
-            hidden_size=config.hidden_size,
-            eps=config.rms_norm_eps,
-            dtype=dtype,
-        )
-        self.post_attention_layernorm = RMSNorm(
-            hidden_size=config.hidden_size,
-            eps=config.rms_norm_eps,
-            dtype=dtype,
-        )
-
-        # Fused QKV A-projection lives inside self_attn namespace
-        # for weight loading compatibility with original checkpoint paths
-        from vllm.model_executor.models.deepseek_v2 import (
-            DeepSeekV2FusedQkvAProjLinear,
-        )
-
-        self.self_attn = nn.Module()
-        self.self_attn.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
-            config.hidden_size,
-            [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
-            quant_config=quant_config,
-            prefix=f"{prefix}.self_attn.fused_qkv_a_proj",
-        )
-
-        # MLA Attention
-        self.attn = DeepseekV32MLAAttention(
-            vllm_config=vllm_config,
-            config=config,
-            hidden_size=config.hidden_size,
-            num_heads=config.num_attention_heads,
-            qk_nope_head_dim=config.qk_nope_head_dim,
-            qk_rope_head_dim=self.qk_rope_head_dim,
-            v_head_dim=config.v_head_dim,
-            q_lora_rank=self.q_lora_rank,
-            kv_lora_rank=self.kv_lora_rank,
-            max_position_embeddings=getattr(config, "max_position_embeddings", 8192),
-            cache_config=cache_config,
-            quant_config=quant_config,
-            topk_indices_buffer=topk_indices_buffer,
-            prefix=f"{prefix}.self_attn",
-        )
-
-        # MoE or Dense MLP
-        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
-        self.is_moe = (
-            config.n_routed_experts is not None
-            and layer_idx >= config.first_k_dense_replace
-            and layer_idx % moe_layer_freq == 0
-        )
-        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
-
-        from vllm.model_executor.models.deepseek_v2 import (
-            DeepseekV2MLP,
-            DeepseekV2MoE,
-        )
-
-        if self.is_moe:
-            self.mlp = DeepseekV2MoE(
-                config=config,
-                parallel_config=parallel_config,
-                quant_config=quant_config,
-                prefix=f"{prefix}.mlp",
-            )
-        else:
-            self.mlp = DeepseekV2MLP(
-                hidden_size=config.hidden_size,
-                intermediate_size=config.intermediate_size,
-                hidden_act=config.hidden_act,
-                quant_config=quant_config,
-                prefix=f"{prefix}.mlp",
-            )
-
-    def forward(
-        self,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        residual: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        if residual is None:
-            residual = hidden_states
-            hidden_states = self.input_layernorm(hidden_states)
-        else:
-            hidden_states, residual = self.input_layernorm(hidden_states, residual)
-
-        # Step 1. hidden_states -> q_c, kv_c, k_pe, index_k, index_weights
-        step1_out = torch.mm(hidden_states, self._fused_step1_hidden_w.T)
-        q_c, kv_c, k_pe, index_k, index_weights = step1_out.split(
-            self._step1_split_sizes,
-            dim=-1,
-        )
-
-        # Steps 2-6. Combined: fused norm/rope + Q projections + sparse MLA.
-        mla = self.attn.mla_attn
-        output_shape = (hidden_states.shape[0], mla.num_heads * mla.v_head_dim)
-        output_dtype = mla.W_UV.dtype
-        attn_out = torch.empty(
-            output_shape,
-            dtype=output_dtype,
-            device=hidden_states.device,
-        )
-        attn_out = torch.ops.vllm.monolithic_attn(
-            positions,
-            q_c,
-            kv_c,
-            k_pe,
-            index_k,
-            index_weights,
-            attn_out,
-            self.layer_name,
-        )
-
-        hidden_states, _ = self.attn.o_proj(attn_out)
-        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
-        hidden_states = self.mlp(hidden_states)
-        return hidden_states, residual
-
-    def fuse_indexer_weights(self) -> None:
-        """Fuse Step 1 and Step 3 BF16 linears used by the inlined path.
-
-        Call after model weights are loaded.
-        """
-        attn = self.attn
-        qkv_a = self.self_attn.fused_qkv_a_proj.weight.data  # [2112, 7168]
-        wk = attn.indexer_wk.weight.data  # [128, 7168]
-        wp = attn.indexer_weights_proj.weight.data  # [64, 7168]
-        if not (qkv_a.dtype == wk.dtype == wp.dtype):
-            raise ValueError(
-                "Cannot fuse Step 1 weights: expected matching dtypes for "
-                "fused_qkv_a_proj, indexer_wk, and indexer_weights_proj."
-            )
-        self._fused_step1_hidden_w = nn.Parameter(
-            torch.cat([qkv_a, wk, wp], dim=0),  # [2304, 7168]
-            requires_grad=False,
-        )
-        self._step1_split_sizes = [
-            self.q_lora_rank,
-            self.kv_lora_rank,
-            self.qk_rope_head_dim,
-            wk.shape[0],
-            wp.shape[0],
-        ]
-
-        wq_b = attn.indexer_wq_b.weight.data
-        q_b = attn.q_b_proj.weight.data
-        if wq_b.dtype != q_b.dtype:
-            raise ValueError(
-                "Cannot fuse Step 3 weights: expected matching dtypes for "
-                "indexer_wq_b and q_b_proj."
-            )
-        self._fused_step3_q_w = nn.Parameter(
-            torch.cat([wq_b, q_b], dim=0),
-            requires_grad=False,
-        )
-        self._q_split_sizes = [wq_b.shape[0], q_b.shape[0]]
-
-
-class DeepseekV32MLAAttention(nn.Module):
-    """
-    MLA attention for DeepSeek V3.2 targeting SM100.
-    MLA forward fully inlined. MLAAttention kept only for KV cache
-    registration and backend/impl initialization.
-    """
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        config,
-        hidden_size: int,
-        num_heads: int,
-        qk_nope_head_dim: int,
-        qk_rope_head_dim: int,
-        v_head_dim: int,
-        q_lora_rank: int,
-        kv_lora_rank: int,
-        max_position_embeddings: int,
-        cache_config: CacheConfig,
-        quant_config: QuantizationConfig | None,
-        topk_indices_buffer: torch.Tensor,
-        prefix: str = "",
-    ) -> None:
-        super().__init__()
-        self.hidden_size = hidden_size
-        self.qk_nope_head_dim = qk_nope_head_dim
-        self.qk_rope_head_dim = qk_rope_head_dim
-        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
-        self.v_head_dim = v_head_dim
-        self.q_lora_rank = q_lora_rank
-        self.kv_lora_rank = kv_lora_rank
-        self.num_heads = num_heads
-        self.num_local_heads = num_heads // get_tensor_model_parallel_world_size()
-        self.scaling = self.qk_head_dim**-0.5
-        self.rms_norm_eps = config.rms_norm_eps
-
-        # Q path
-        self.q_a_layernorm_weight = nn.Parameter(
-            torch.ones(q_lora_rank, dtype=torch.get_default_dtype())
-        )
-        self.q_b_proj = ColumnParallelLinear(
-            q_lora_rank,
-            num_heads * self.qk_head_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.q_b_proj",
-        )
-
-        # KV path
-        self.kv_a_layernorm_weight = nn.Parameter(
-            torch.ones(kv_lora_rank, dtype=torch.get_default_dtype())
-        )
-        self.kv_b_proj = ColumnParallelLinear(
-            kv_lora_rank,
-            num_heads * (qk_nope_head_dim + v_head_dim),
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.kv_b_proj",
-        )
-
-        # Output projection (TP sync point)
-        self.o_proj = RowParallelLinear(
-            num_heads * v_head_dim,
-            hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj",
-        )
-
-        # RoPE
-        if config.rope_parameters["rope_type"] != "default":
-            config.rope_parameters["rope_type"] = (
-                "deepseek_yarn"
-                if config.rope_parameters.get("apply_yarn_scaling", True)
-                else "deepseek_llama_scaling"
-            )
-        self.rotary_emb = get_rope(
-            qk_rope_head_dim,
-            max_position=max_position_embeddings,
-            rope_parameters=config.rope_parameters,
-            is_neox_style=False,
-        )
-        if config.rope_parameters["rope_type"] == "deepseek_yarn":
-            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
-            scaling_factor = config.rope_parameters["factor"]
-            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
-            self.scaling = self.scaling * mscale * mscale
-
-        # V3.2 Sparse Indexer (inlined)
-        self.indexer_rope_emb = get_rope(
-            qk_rope_head_dim,
-            max_position=max_position_embeddings,
-            rope_parameters=config.rope_parameters,
-            is_neox_style=not getattr(config, "indexer_rope_interleave", False),
-        )
-        self.topk_tokens = config.index_topk
-        self.index_n_heads = config.index_n_heads
-        self.index_head_dim = config.index_head_dim
-        self.indexer_softmax_scale = config.index_head_dim**-0.5
-        self.indexer_quant_block_size = 128
-        self.topk_indices_buffer = topk_indices_buffer
-
-        self.indexer_wq_b = ReplicatedLinear(
-            q_lora_rank,
-            config.index_head_dim * config.index_n_heads,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.indexer.wq_b",
-        )
-        self.indexer_wk = ReplicatedLinear(
-            hidden_size,
-            config.index_head_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.indexer.wk",
-        )
-        self.indexer_k_norm = LayerNorm(config.index_head_dim, eps=1e-6)
-        self.indexer_weights_proj = ReplicatedLinear(
-            hidden_size,
-            config.index_n_heads,
-            bias=False,
-            quant_config=None,
-            prefix=f"{prefix}.indexer.weights_proj",
-        )
-
-        idx_dim = config.index_head_dim
-        indexer_cache_head_dim = idx_dim + idx_dim // 128 * 4
-        self.indexer_k_cache = DeepseekV32IndexerCache(
-            head_dim=indexer_cache_head_dim,
-            dtype=torch.uint8,
-            prefix=f"{prefix}.indexer.k_cache",
-            cache_config=cache_config,
-        )
-        self.indexer_op = SparseAttnIndexer(
-            self.indexer_k_cache,
-            self.indexer_quant_block_size,
-            "ue8m0",
-            self.topk_tokens,
-            config.index_head_dim,
-            vllm_config.model_config.max_model_len,
-            get_max_prefill_buffer_size(vllm_config),
-            self.topk_indices_buffer,
-        )
-
-        # MLAAttention stub: only for KV cache registration + backend init.
-        # We never call its forward(); we inline everything below.
-        class _IndexerProxy:
-            def __init__(proxy_self):
-                proxy_self.topk_indices_buffer = topk_indices_buffer
-                proxy_self.indexer_op = self.indexer_op
-
-        self._indexer_proxy = _IndexerProxy()
-        self.mla_attn = MLAAttention(
-            num_heads=self.num_local_heads,
-            scale=self.scaling,
-            qk_nope_head_dim=qk_nope_head_dim,
-            qk_rope_head_dim=qk_rope_head_dim,
-            v_head_dim=v_head_dim,
-            q_lora_rank=q_lora_rank,
-            kv_lora_rank=kv_lora_rank,
-            kv_b_proj=self.kv_b_proj,
-            cache_config=cache_config,
-            quant_config=quant_config,
-            prefix=f"{prefix}.mla_attn",
-            use_sparse=True,
-            indexer=self._indexer_proxy,
-        )
-
-
-def remap_weight_name(name: str) -> str:
-    """Remap checkpoint names that differ from the module layout."""
-    replacements = [
-        (
-            "self_attn.q_a_layernorm.weight",
-            "attn.q_a_layernorm_weight",
-        ),
-        (
-            "self_attn.kv_a_layernorm.weight",
-            "attn.kv_a_layernorm_weight",
-        ),
-        ("self_attn.q_b_proj", "attn.q_b_proj"),
-        ("self_attn.kv_b_proj", "attn.kv_b_proj"),
-        ("self_attn.o_proj", "attn.o_proj"),
-        ("self_attn.indexer.", "attn.indexer_"),
-    ]
-    for old, new in replacements:
-        if old in name:
-            return name.replace(old, new)
-    return name

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/model.py

@@ -1,151 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""DeepSeek V3.2 NVFP4 model for SM100 (Blackwell)."""
-
-from collections.abc import Iterable
-
-import torch
-from torch import nn
-
-from vllm.compilation.decorators import support_torch_compile
-from vllm.config import VllmConfig
-from vllm.distributed import get_tensor_model_parallel_world_size
-from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead,
-    VocabParallelEmbedding,
-)
-from vllm.platforms import current_platform
-
-from .layer import DeepseekV32DecoderLayer, remap_weight_name
-
-logger = init_logger(__name__)
-
-
-@support_torch_compile
-class DeepseekV32Model(nn.Module):
-    fall_back_to_pt_during_load = False
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        super().__init__()
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-        self.config = config
-        self.device = current_platform.device_type
-
-        topk_tokens = config.index_topk
-        self.topk_indices_buffer = torch.empty(
-            vllm_config.scheduler_config.max_num_batched_tokens,
-            topk_tokens,
-            dtype=torch.int32,
-            device=self.device,
-        )
-
-        self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size,
-            config.hidden_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.embed_tokens",
-        )
-
-        self.layers = nn.ModuleList(
-            [
-                DeepseekV32DecoderLayer(
-                    vllm_config=vllm_config,
-                    config=config,
-                    layer_idx=i,
-                    topk_indices_buffer=self.topk_indices_buffer,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(config.num_hidden_layers)
-            ]
-        )
-
-        self.norm = RMSNorm(
-            hidden_size=config.hidden_size,
-            eps=config.rms_norm_eps,
-            dtype=torch.get_default_dtype(),
-        )
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-    ) -> torch.Tensor:
-        hidden_states = self.embed_tokens(input_ids)
-        residual = None
-        for layer in self.layers:
-            hidden_states, residual = layer(positions, hidden_states, residual)
-        hidden_states, _ = self.norm(hidden_states, residual)
-        return hidden_states
-
-
-class DeepseekV32ForCausalLM(nn.Module):
-    packed_modules_mapping = {
-        "gate_up_proj": ["gate_proj", "up_proj"],
-        "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
-    }
-
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
-        super().__init__()
-        config = vllm_config.model_config.hf_config
-        quant_config = vllm_config.quant_config
-        self.config = config
-        self.quant_config = quant_config
-        self.tp_size = get_tensor_model_parallel_world_size()
-
-        self.model = DeepseekV32Model(
-            vllm_config=vllm_config,
-            prefix=f"{prefix}.model" if prefix else "model",
-        )
-        self.lm_head = ParallelLMHead(
-            config.vocab_size,
-            config.hidden_size,
-            quant_config=quant_config,
-            prefix=f"{prefix}.lm_head" if prefix else "lm_head",
-        )
-        self.logits_processor = LogitsProcessor(config.vocab_size)
-        self.num_redundant_experts = 0
-
-    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.model.embed_tokens(input_ids)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        intermediate_tensors=None,
-        inputs_embeds=None,
-    ) -> torch.Tensor:
-        return self.model(input_ids, positions)
-
-    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
-        return self.logits_processor(self.lm_head, hidden_states)
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Delegate to the original DeepSeek V2 weight loader.
-
-        Our module structure matches the original for all weights that
-        need special loading (fused_qkv_a_proj, experts, gate_up_proj).
-        Only layernorm weights and indexer paths differ.
-        """
-        from vllm.model_executor.models.deepseek_v2 import (
-            DeepseekV2ForCausalLM,
-        )
-
-        def _remap_weights():
-            for name, w in weights:
-                yield remap_weight_name(name), w
-
-        self.use_mha = False
-        self.fuse_qkv_a_proj = True
-        self.is_fp4_ckpt = False
-        loaded = DeepseekV2ForCausalLM.load_weights(self, _remap_weights())
-
-        # Fuse indexer linear weights after loading.
-        for layer in self.model.layers:
-            layer.fuse_indexer_weights()
-
-        return loaded

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/mtp.py

@@ -1,209 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""DeepSeek V3.2 MTP model for SM100 (Blackwell)."""
-
-from collections.abc import Iterable
-
-import torch
-import torch.nn as nn
-
-from vllm.compilation.decorators import support_torch_compile
-from vllm.config import VllmConfig
-from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    VocabParallelEmbedding,
-)
-from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepSeekMTPBase
-from vllm.model_executor.models.deepseek_mtp import (
-    DeepSeekMultiTokenPredictor as DeepSeekMultiTokenPredictorBase,
-)
-from vllm.model_executor.models.deepseek_mtp import (
-    DeepSeekMultiTokenPredictorLayer as DeepSeekMultiTokenPredictorLayerBase,
-)
-from vllm.model_executor.models.deepseek_mtp import SharedHead as SharedHeadBase
-from vllm.model_executor.models.deepseek_v2 import DeepseekV2MoE
-from vllm.model_executor.models.utils import maybe_prefix
-from vllm.platforms import current_platform
-from vllm.sequence import IntermediateTensors
-
-from .kernels import fused_mtp_entry
-from .layer import DeepseekV32DecoderLayer
-from .model import remap_weight_name
-
-logger = init_logger(__name__)
-
-
-class SharedHead(SharedHeadBase):
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        return rms_norm(hidden_states, self.norm.weight, self.norm.variance_epsilon)
-
-
-class DeepSeekMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayerBase):
-    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
-        nn.Module.__init__(self)
-
-        assert vllm_config.speculative_config is not None
-        config = vllm_config.speculative_config.draft_model_config.hf_config
-        quant_config = vllm_config.quant_config
-
-        self.config = config
-        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
-
-        topk_indices_buffer = torch.empty(
-            vllm_config.scheduler_config.max_num_batched_tokens,
-            config.index_topk,
-            dtype=torch.int32,
-            device=current_platform.device_type,
-        )
-
-        self.shared_head = SharedHead(
-            config=config, prefix=prefix, quant_config=quant_config
-        )
-        self.mtp_block = DeepseekV32DecoderLayer(
-            vllm_config=vllm_config,
-            config=config,
-            layer_idx=int(prefix.rsplit(".", 1)[-1]),
-            topk_indices_buffer=topk_indices_buffer,
-            prefix=prefix,
-        )
-        # Pre-allocated 0-dim eps tensors so fused_mtp_entry can stay
-        # tensor-only (avoids Python-float scalars leaking into the
-        # torch.compile input list).
-        self._e_eps_gpu = torch.full(
-            (),
-            self.enorm.variance_epsilon,
-            dtype=torch.float32,
-            device=current_platform.device_type,
-        )
-        self._h_eps_gpu = torch.full(
-            (),
-            self.hnorm.variance_epsilon,
-            dtype=torch.float32,
-            device=current_platform.device_type,
-        )
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        positions: torch.Tensor,
-        previous_hidden_states: torch.Tensor,
-        inputs_embeds: torch.Tensor | None = None,
-        spec_step_index: int = 0,
-    ) -> torch.Tensor:
-        assert inputs_embeds is not None
-        eh_concat = fused_mtp_entry(
-            inputs_embeds,
-            previous_hidden_states,
-            positions,
-            self.enorm.weight,
-            self.hnorm.weight,
-            self._e_eps_gpu,
-            self._h_eps_gpu,
-        )
-        hidden_states = self.eh_proj(eh_concat)
-        hidden_states, residual = self.mtp_block(
-            positions=positions, hidden_states=hidden_states, residual=None
-        )
-        hidden_states = residual + hidden_states
-        return hidden_states
-
-
-class DeepSeekMultiTokenPredictor(DeepSeekMultiTokenPredictorBase):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        nn.Module.__init__(self)
-        config = vllm_config.model_config.hf_config
-        self.mtp_start_layer_idx = config.num_hidden_layers
-        self.num_mtp_layers = config.num_nextn_predict_layers
-
-        self.layers = torch.nn.ModuleDict(
-            {
-                str(idx): DeepSeekMultiTokenPredictorLayer(
-                    vllm_config, f"{prefix}.layers.{idx}"
-                )
-                for idx in range(
-                    self.mtp_start_layer_idx,
-                    self.mtp_start_layer_idx + self.num_mtp_layers,
-                )
-            }
-        )
-        self.embed_tokens = VocabParallelEmbedding(
-            config.vocab_size,
-            config.hidden_size,
-            prefix=maybe_prefix(prefix, "embed_tokens"),
-        )
-        self.logits_processor = LogitsProcessor(config.vocab_size)
-
-
-@support_torch_compile
-class DeepSeekMTP(DeepSeekMTPBase):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
-        nn.Module.__init__(self)
-        self.config = vllm_config.model_config.hf_config
-        self.quant_config = vllm_config.quant_config
-        assert hasattr(self.config, "index_topk")
-        cache_config = vllm_config.cache_config
-        if cache_config.cache_dtype == "bfloat16":
-            cache_config.cache_dtype = "auto"
-            logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
-        self.model = DeepSeekMultiTokenPredictor(
-            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
-        )
-        self.set_moe_parameters()
-        # Keep the original loader from applying the fused FP4 indexer remap.
-        self.is_fp4_ckpt = False
-
-    def set_moe_parameters(self):
-        self.expert_weights = []
-        self.num_moe_layers = self.config.num_nextn_predict_layers
-        self.num_expert_groups = self.config.n_group
-
-        self.moe_layers = []
-        self.moe_mlp_layers = []
-        example_moe = None
-        for layer in self.model.layers.values():
-            layer = layer.mtp_block
-            assert isinstance(layer, DeepseekV32DecoderLayer)
-            if isinstance(layer.mlp, DeepseekV2MoE):
-                example_moe = layer.mlp
-                self.moe_mlp_layers.append(layer.mlp)
-                self.moe_layers.append(layer.mlp.experts)
-        self.extract_moe_parameters(example_moe)
-
-    def forward(
-        self,
-        input_ids: torch.Tensor | None,
-        positions: torch.Tensor,
-        hidden_states: torch.Tensor,
-        intermediate_tensors: IntermediateTensors | None = None,
-        inputs_embeds: torch.Tensor | None = None,
-        spec_step_idx: int = 0,
-    ) -> torch.Tensor:
-        del intermediate_tensors
-        return self.model(
-            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
-        )
-
-    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        loaded_params = super().load_weights(weights)
-        for layer in self.model.layers.values():
-            layer.mtp_block.fuse_indexer_weights()
-        return loaded_params
-
-    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
-        name = super()._rewrite_spec_layer_name(spec_layer, name)
-        return remap_weight_name(name)
-
-
-@torch.compile
-def rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
-    orig_dtype = x.dtype
-    x = x.to(torch.float32)
-    mean_sq = (x * x).mean(dim=-1, keepdim=True)
-    rrms = torch.rsqrt(mean_sq + eps)
-    x = x * rrms
-    x = x * w.to(torch.float32)
-    return x.to(orig_dtype)

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/sparse_indexer.py

@@ -1,175 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Custom Sparse Attention Indexer layers."""
-
-import torch
-
-import vllm.envs as envs
-from vllm import _custom_ops as ops
-from vllm.forward_context import get_forward_context
-from vllm.logger import init_logger
-from vllm.platforms import current_platform
-from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
-from vllm.utils.torch_utils import (
-    LayerNameType,
-    _resolve_layer_name,
-)
-from vllm.v1.attention.backends.mla.indexer import (
-    DeepseekV32IndexerMetadata,
-)
-from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
-from vllm.v1.worker.workspace import current_workspace_manager
-
-logger = init_logger(__name__)
-
-RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
-
-
-def sparse_attn_indexer(
-    k_cache_prefix: LayerNameType,
-    kv_cache: torch.Tensor,
-    q_fp8: torch.Tensor,
-    weights: torch.Tensor,
-    topk_tokens: int,
-    head_dim: int,
-    max_model_len: int,
-    total_seq_lens: int,
-    topk_indices_buffer: torch.Tensor,
-) -> torch.Tensor:
-    # careful! this will be None in dummy run
-    attn_metadata = get_forward_context().attn_metadata
-    fp8_dtype = current_platform.fp8_dtype()
-    k_cache_prefix = _resolve_layer_name(k_cache_prefix)
-
-    # assert isinstance(attn_metadata, dict)
-    if not isinstance(attn_metadata, dict):
-        # Reserve workspace for indexer during profiling run
-        current_workspace_manager().get_simultaneous(
-            ((total_seq_lens, head_dim), torch.float8_e4m3fn),
-            ((total_seq_lens, 4), torch.uint8),
-            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-        )
-
-        # Dummy allocation to simulate for peak logits tensor memory during inference.
-        # FP8 elements so elements == bytes
-        max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
-        _ = torch.empty(max_logits_elems, dtype=torch.uint8, device=q_fp8.device)
-        return None
-
-    attn_metadata = attn_metadata[k_cache_prefix]  # type: ignore[assignment]
-    assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
-    has_decode = attn_metadata.num_decodes > 0
-    has_prefill = attn_metadata.num_prefills > 0
-    num_decode_tokens = attn_metadata.num_decode_tokens
-
-    if has_prefill:
-        prefill_metadata = attn_metadata.prefill
-        assert prefill_metadata is not None
-
-        # Get the full shared workspace buffers once (will allocate on first use)
-        workspace_manager = current_workspace_manager()
-        k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
-            ((total_seq_lens, head_dim), fp8_dtype),
-            ((total_seq_lens, 4), torch.uint8),
-        )
-        for chunk in prefill_metadata.chunks:
-            k_fp8 = k_fp8_full[: chunk.total_seq_lens]
-            k_scale = k_scale_full[: chunk.total_seq_lens]
-
-            if not chunk.skip_kv_gather:
-                ops.cp_gather_indexer_k_quant_cache(
-                    kv_cache,
-                    k_fp8,
-                    k_scale,
-                    chunk.block_table,
-                    chunk.cu_seq_lens,
-                )
-
-            logits = fp8_mqa_logits(
-                q_fp8[chunk.token_start : chunk.token_end],
-                (k_fp8, k_scale.view(torch.float32).flatten()),
-                weights[chunk.token_start : chunk.token_end],
-                chunk.cu_seqlen_ks,
-                chunk.cu_seqlen_ke,
-                clean_logits=False,
-            )
-            num_rows = logits.shape[0]
-
-            topk_indices = topk_indices_buffer[
-                chunk.token_start : chunk.token_end, :topk_tokens
-            ]
-
-            torch.ops._C.top_k_per_row_prefill(
-                logits,
-                chunk.cu_seqlen_ks,
-                chunk.cu_seqlen_ke,
-                topk_indices,
-                num_rows,
-                logits.stride(0),
-                logits.stride(1),
-                topk_tokens,
-            )
-
-    if has_decode:
-        decode_metadata = attn_metadata.decode
-        assert decode_metadata is not None
-        # kv_cache shape [
-        # kv_cache size requirement [num_block, block_size, n_head, head_dim],
-        # we only have [num_block, block_size, head_dim],
-        kv_cache = kv_cache.unsqueeze(-2)
-        decode_lens = decode_metadata.decode_lens
-        if decode_metadata.requires_padding:
-            # pad in edge case where we have short chunked prefill length <
-            # decode_threshold since we unstrictly split
-            # prefill and decode by decode_threshold
-            # (currently set to 1 + speculative tokens)
-            padded_q_fp8_decode_tokens = pack_seq_triton(
-                q_fp8[:num_decode_tokens], decode_lens
-            )
-        else:
-            padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
-                decode_lens.shape[0], -1, *q_fp8.shape[1:]
-            )
-        # TODO: move and optimize below logic with triton kernels
-        batch_size = padded_q_fp8_decode_tokens.shape[0]
-        next_n = padded_q_fp8_decode_tokens.shape[1]
-        num_padded_tokens = batch_size * next_n
-        seq_lens = decode_metadata.seq_lens[:batch_size]
-        # seq_lens is (B, next_n) for native spec decode, (B,) otherwise.
-        # fp8_paged_mqa_logits and all topk kernels accept both shapes.
-        logits = fp8_paged_mqa_logits(
-            padded_q_fp8_decode_tokens,
-            kv_cache,
-            weights[:num_padded_tokens],
-            seq_lens,
-            decode_metadata.block_table,
-            decode_metadata.schedule_metadata,
-            max_model_len=max_model_len,
-            clean_logits=False,
-        )
-        num_rows = logits.shape[0]
-        topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
-
-        workspace_manager = current_workspace_manager()
-        (topk_workspace,) = workspace_manager.get_simultaneous(
-            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
-        )
-        torch.ops._C.persistent_topk(
-            logits,
-            seq_lens,
-            topk_indices,
-            topk_workspace,
-            topk_tokens,
-            attn_metadata.max_seq_len,
-        )
-
-        if decode_metadata.requires_padding:
-            # if padded, we need to unpack
-            # the topk indices removing padded tokens
-            topk_indices = unpack_seq_triton(
-                topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
-                decode_lens,
-            )
-            topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = (
-                topk_indices
-            )

vllm/platforms/cpu.py

@@ -6,15 +6,16 @@ import os
 import platform
 import subprocess
 import sys
-from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-import psutil
 import torch
 
-from vllm import envs
 from vllm.logger import init_logger
-from vllm.utils.ompmultiprocessing import OMPProcessManager
+from vllm.utils.cpu_resource_utils import (
+    DEVICE_CONTROL_ENV_VAR,
+    get_memory_node_info,
+)
+from vllm.utils.mem_constants import GiB_bytes
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -38,49 +39,13 @@ def get_max_threads(pid=0):
         raise NotImplementedError("Unsupported OS")
 
 
-@dataclass
-class LogicalCPUInfo:
-    id: int = -1
-    physical_core: int = -1
-    numa_node: int = -1
-
-    @classmethod
-    def _int(cls, value: str) -> int:
-        try:
-            int_value = int(value)
-        except Exception:
-            int_value = -1
-        return int_value
-
-    @staticmethod
-    def json_decoder(obj_dict: dict):
-        id = obj_dict.get("cpu")
-        physical_core = obj_dict.get("core")
-        numa_node = obj_dict.get("node")
-
-        if not (id is None or physical_core is None or numa_node is None):
-            return LogicalCPUInfo(
-                id=LogicalCPUInfo._int(id),
-                physical_core=LogicalCPUInfo._int(physical_core),
-                numa_node=LogicalCPUInfo._int(numa_node),
-            )
-        else:
-            return obj_dict
-
-
 class CpuPlatform(Platform):
     _enum = PlatformEnum.CPU
     device_name: str = "cpu"
     device_type: str = "cpu"
     dispatch_key: str = "CPU"
     dist_backend: str = "gloo"
-    device_control_env_var = "CPU_VISIBLE_MEMORY_NODES"
-    omp_process_manager = None
-    # Simultaneous Multithreading (SMT) level for OpenMP:
-    # 4 on PowerPC, 1 on non-PowerPC architectures
-    smt = 1
-    global_cpu_mask = None
-    simulate_numa = int(os.environ.get("_SIM_MULTI_NUMA", 0))
+    device_control_env_var = DEVICE_CONTROL_ENV_VAR
 
     @property
     def supported_dtypes(self) -> list[torch.dtype]:
@@ -123,29 +88,9 @@ class CpuPlatform(Platform):
 
     @classmethod
     def get_device_total_memory(cls, device_id: int = 0) -> int:
-        from vllm.utils.mem_constants import GiB_bytes
-        from vllm.utils.mem_utils import format_gib
+        meminfo = get_memory_node_info(device_id)
 
-        kv_cache_space = envs.VLLM_CPU_KVCACHE_SPACE
-        node_dir = "/sys/devices/system/node"
-        if kv_cache_space is None:
-            nodes = (
-                [d for d in os.listdir(node_dir) if d.startswith("node")]
-                if os.path.exists(node_dir)
-                else []
-            )
-            num_numa_nodes = len(nodes) or 1
-            free_cpu_memory = psutil.virtual_memory().total // num_numa_nodes
-            DEFAULT_CPU_MEM_UTILIZATION = 0.5
-            kv_cache_space = int(free_cpu_memory * DEFAULT_CPU_MEM_UTILIZATION)
-            logger.warning_once(
-                "VLLM_CPU_KVCACHE_SPACE not set. Using %s GiB for KV cache.",
-                format_gib(kv_cache_space),
-            )
-        else:
-            kv_cache_space *= GiB_bytes
-
-        return kv_cache_space
+        return meminfo.total_memory
 
     @classmethod
     def set_device(cls, device: torch.device) -> None:
@@ -180,6 +125,12 @@ class CpuPlatform(Platform):
                 "otherwise the performance is not optimized."
             )
 
+        # Lagecy setting
+        env_key = "VLLM_CPU_KVCACHE_SPACE"
+        if env_key in os.environ and os.environ[env_key] != "":
+            kv_cache_space = int(os.environ[env_key])
+            cache_config.kv_cache_memory_bytes = kv_cache_space * GiB_bytes
+
         scheduler_config = vllm_config.scheduler_config
         # async scheduling is not required on CPU
         scheduler_config.async_scheduling = False
@@ -198,8 +149,6 @@ class CpuPlatform(Platform):
             )
             cache_config.cache_dtype = "auto"
 
-        cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory()
-
         parallel_config = vllm_config.parallel_config
         # OMP requires the MP executor to function correctly, UniProc is not
         # supported as it is not possible to set the OMP environment correctly
@@ -278,21 +227,45 @@ class CpuPlatform(Platform):
         os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"
 
         ld_preload_str = os.getenv("LD_PRELOAD", "")
-
-        # Intel and CLANG OpenMP setting
-        if "libiomp5.so" in ld_preload_str or "libomp5" in ld_preload_str:
-            # The time(milliseconds) that a thread should wait after
-            # completing the execution of a parallel region, before sleeping.
-            os.environ["KMP_BLOCKTIME"] = "1"
-            # Prevents the CPU to run into low performance state
-            os.environ["KMP_TPAUSE"] = "0"
-            # Provides fine granularity parallelism
-            os.environ["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist"
-            os.environ["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist"
-            os.environ["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist"
-
         cpu_architecture = Platform.get_cpu_architecture()
 
+        if (
+            platform.system() == "Linux"
+            and cpu_architecture
+            in (CpuArchEnum.ARM, CpuArchEnum.POWERPC, CpuArchEnum.X86)
+            and not (
+                "libomp" in ld_preload_str
+                or "libgomp" in ld_preload_str
+                or "libiomp" in ld_preload_str
+            )
+        ):
+            # We need to LD_PRELOAD PyTorch's libgomp, otherwise only
+            # one core will be properly utilized when we thread-bind
+            # See: https://github.com/vllm-project/vllm/issues/27369
+            # TODO: Remove once:
+            # https://github.com/pytorch/pytorch/issues/166087 is fixed
+
+            # We need to find the location of PyTorch's libgomp
+            torch_pkg = os.path.dirname(torch.__file__)
+            site_root = os.path.dirname(torch_pkg)
+            # Search both torch.libs and torch/lib - See:
+            # https://github.com/vllm-project/vllm/issues/30470
+            torch_libs_paths = [
+                os.path.join(site_root, "torch.libs"),
+                os.path.join(torch_pkg, "lib"),
+            ]
+            pytorch_libgomp_so_candidates = []
+            for torch_libs in torch_libs_paths:
+                pytorch_libgomp_so_candidates.extend(
+                    glob.glob(os.path.join(torch_libs, "libgomp*.so*"))
+                )
+            if pytorch_libgomp_so_candidates:
+                pytorch_libgomp_so = pytorch_libgomp_so_candidates[0]
+                if ld_preload_str:
+                    ld_preload_str += ":"
+                ld_preload_str += pytorch_libgomp_so
+                os.environ["LD_PRELOAD"] = ld_preload_str
+
         # LD_PRELOAD libtcmalloc, bundled under vllm/libs to reduce
         # memory allocation overhead
         if (
@@ -331,13 +304,6 @@ class CpuPlatform(Platform):
                 vllm_config.model_config.max_model_len,
                 vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,
             )
-        # CI specific "quick" NUMA simulation - split all available CPUs
-        # into a fake NUMA topology
-        if os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", None) is not None:
-            os.environ["_SIM_MULTI_NUMA"] = str(
-                vllm_config.parallel_config.world_size
-                * vllm_config.parallel_config._api_process_count
-            )
 
     @classmethod
     def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
@@ -345,78 +311,6 @@ class CpuPlatform(Platform):
         # Move that logic here so block_size is chosen by the backend.
         pass
 
-    @classmethod
-    def get_omp_manager(cls) -> OMPProcessManager:
-        # initialise the OMP resource management if need be and return the manager
-        if cls.omp_process_manager is None:
-            if cls.get_cpu_architecture() == CpuArchEnum.POWERPC:
-                cls.smt = 4
-            cls.omp_process_manager = OMPProcessManager(
-                affinity=cls.get_global_cpu_mask(), smt=cls.smt
-            )
-            # we need to fix up the topology returned by the OMP Manager for
-            # simulated NUMA environments in CI
-            if cls.simulate_numa > 0:
-                logger.info(
-                    "Adjusting numa topology to resemble at least %d nodes",
-                    int(cls.simulate_numa),
-                )
-                om = cls.omp_process_manager
-                while len(om.omp_places) < cls.simulate_numa:
-                    new_omp_places = []
-                    touched = False
-                    for omp_place in om.omp_places:
-                        if len(omp_place["mask"]) > 1:
-                            touched = True
-                            cpu_list = sorted(list(omp_place["mask"]))
-                            new_omp_places.append(
-                                {
-                                    "mask": set(cpu_list[0 : int(len(cpu_list) / 2)]),
-                                    "available": True,
-                                }
-                            )
-                            new_omp_places.append(
-                                {
-                                    "mask": set(cpu_list[int(len(cpu_list) / 2) :]),
-                                    "available": True,
-                                }
-                            )
-                    if touched:
-                        om.omp_places = new_omp_places
-                    else:
-                        raise ValueError(
-                            "Cannot split the existing NUMA topology to match "
-                            "simulation requirements"
-                        )
-
-        return cls.omp_process_manager
-
-    @classmethod
-    def get_global_cpu_mask(cls) -> set[int]:
-        # get global cpu mask
-        if cls.global_cpu_mask is None:
-            if hasattr(os, "sched_getaffinity"):
-                cls.global_cpu_mask = os.sched_getaffinity(0)
-            else:
-                # macOS does not support sched_getaffinity
-                cpu_count = os.cpu_count() or 1
-                cls.global_cpu_mask = set(range(cpu_count))
-        return cls.global_cpu_mask
-
-    @classmethod
-    def reserve_cpus(cls, reserve: set[int]) -> bool:
-        # remove CPUs from global mask, for now there is no "release" mechanism
-        if cls.omp_process_manager is not None:
-            for place in cls.omp_process_manager.omp_places:
-                if not place["available"]:
-                    return False
-        cls.global_cpu_mask = cls.get_global_cpu_mask() - reserve
-        # reinitialize OMP resource management
-        cls.omp_process_manager = OMPProcessManager(
-            affinity=cls.global_cpu_mask, smt=cls.smt
-        )
-        return True
-
     @classmethod
     def discover_numa_topology(cls) -> list[list[int]]:
         """

vllm/platforms/cuda.py

@@ -108,6 +108,22 @@ def _get_backend_priorities(
                         AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
                     ]
 
+            return [
+                AttentionBackendEnum.FLASHINFER_MLA,
+                AttentionBackendEnum.CUTLASS_MLA,
+                AttentionBackendEnum.FLASH_ATTN_MLA,
+                AttentionBackendEnum.FLASHMLA,
+                AttentionBackendEnum.TRITON_MLA,
+                *sparse_backends,
+            ]
+        elif device_capability.major == 9:
+            # Hopper (SM90) — CUTLASS FA3 is highest priority for sparse MLA
+            # with BF16 KV cache. Falls back to FlashMLA Sparse for FP8.
+            sparse_backends = [
+                AttentionBackendEnum.CUTLASS_FA3_MLA_SPARSE,
+                AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
+                AttentionBackendEnum.FLASHMLA_SPARSE,
+            ]
             return [
                 AttentionBackendEnum.FLASHINFER_MLA,
                 AttentionBackendEnum.CUTLASS_MLA,
@@ -131,6 +147,7 @@ def _get_backend_priorities(
                 AttentionBackendEnum.FLASH_ATTN,
                 AttentionBackendEnum.TRITON_ATTN,
                 AttentionBackendEnum.FLEX_ATTENTION,
+                AttentionBackendEnum.TURBOQUANT,
             ]
         else:
             return [
@@ -138,6 +155,7 @@ def _get_backend_priorities(
                 AttentionBackendEnum.FLASHINFER,
                 AttentionBackendEnum.TRITON_ATTN,
                 AttentionBackendEnum.FLEX_ATTENTION,
+                AttentionBackendEnum.TURBOQUANT,
             ]
 
 
@@ -255,11 +273,6 @@ class CudaPlatformBase(Platform):
         valid_backends_priorities = []
         invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
 
-        # TurboQuant KV cache: route directly to TQ backend
-        kv_cache_dtype = attn_selector_config.kv_cache_dtype
-        if kv_cache_dtype is not None and kv_cache_dtype.startswith("turboquant_"):
-            return [(AttentionBackendEnum.TURBOQUANT, 0)], {}
-
         backend_priorities = _get_backend_priorities(
             attn_selector_config.use_mla,
             device_capability,

vllm/tool_parsers/qwen3xml_tool_parser.py

@@ -1258,11 +1258,11 @@ class Qwen3XMLToolParser(ToolParser):
             return None
 
         # Parse the delta text and get the result
-        result = self.parser.parse_single_streaming_chunks(delta_text)
+        delta = self.parser.parse_single_streaming_chunks(delta_text)
 
         # Update tool call tracking arrays based on incremental parsing results
-        if result and result.tool_calls:
-            for tool_call in result.tool_calls:
+        if delta and delta.tool_calls:
+            for tool_call in delta.tool_calls:
                 if tool_call.function:
                     tool_index = (
                         tool_call.index
@@ -1292,4 +1292,7 @@ class Qwen3XMLToolParser(ToolParser):
                         self.streamed_args_for_tool[tool_index] += (
                             tool_call.function.arguments
                         )
-        return result
+        if delta.content is None and not delta.tool_calls and delta.reasoning is None:
+            # If no content and no tool calls, return None to indicate no update
+            return None
+        return delta

vllm/utils/cpu_resource_utils.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import os
+import platform
+import subprocess
+from dataclasses import dataclass
+from functools import cache
+
+import psutil
+import regex as re
+
+DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES"
+
+
+@dataclass
+class LogicalCPUInfo:
+    id: int = -1
+    physical_core: int = -1
+    numa_node: int = -1
+
+    @classmethod
+    def _int(cls, value: str) -> int:
+        try:
+            int_value = int(value)
+        except Exception:
+            int_value = -1
+        return int_value
+
+    @staticmethod
+    def json_decoder(obj_dict: dict):
+        id = obj_dict.get("cpu")
+        physical_core = obj_dict.get("core")
+        numa_node = obj_dict.get("node")
+
+        if not (id is None or physical_core is None or numa_node is None):
+            return LogicalCPUInfo(
+                id=LogicalCPUInfo._int(id),
+                physical_core=LogicalCPUInfo._int(physical_core),
+                numa_node=LogicalCPUInfo._int(numa_node),
+            )
+        else:
+            return obj_dict
+
+
+@dataclass
+class MemoryNodeInfo:
+    total_memory: int = -1
+    available_memory: int = -1
+
+
+def get_memory_affinity(pid: int = 0) -> list[int]:
+    pid = os.getpid() if pid == 0 else pid
+    path = f"/proc/{pid}/status"
+    with open(path) as f:
+        for line in f:
+            if line.startswith("Mems_allowed_list:"):
+                # Extract the string part (e.g., "0-1,3")
+                raw_list = line.split(":")[1].strip()
+                return parse_id_list(raw_list)
+    return []
+
+
+def parse_id_list(raw_str: str) -> list[int]:
+    """Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8]"""
+    result: list[int] = []
+    if not raw_str:
+        return result
+
+    for part in raw_str.split(","):
+        if "-" in part:
+            start, end = map(int, part.split("-"))
+            result.extend(range(start, end + 1))
+        else:
+            result.append(int(part))
+    return sorted(list(set(result)))
+
+
+def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
+    if platform.system() == "Darwin":
+        # MacOS has no memory node
+        return MemoryNodeInfo(
+            total_memory=psutil.virtual_memory().total,
+            available_memory=psutil.virtual_memory().available,
+        )
+
+    meminfo_path = f"/sys/devices/system/node/node{node_id}/meminfo"
+    if not os.path.exists(meminfo_path):
+        raise RuntimeError(f"{meminfo_path} doesn't exit.")
+
+    meminfo = {}
+    with open(meminfo_path) as f:
+        for line in f:
+            # Each line looks like: "Node 0 MemTotal: 97421888 kB"
+            parts = line.split()
+            key = parts[2].rstrip(":")
+            # convert to Bytes
+            value = int(parts[3]) * 1024
+            meminfo[key] = value
+
+    total_memory = meminfo["MemTotal"]
+    free_memory = meminfo["MemFree"]
+    active_file_memory = meminfo["Active(file)"]
+    inactive_file_memory = meminfo["Inactive(file)"]
+    reclaimable_memory = meminfo["SReclaimable"]
+    available_memory = (
+        free_memory + active_file_memory + inactive_file_memory + reclaimable_memory
+    )
+
+    return MemoryNodeInfo(
+        total_memory=total_memory,
+        available_memory=available_memory,
+    )
+
+
+def get_allowed_cpu_list() -> list[LogicalCPUInfo]:
+    cpu_list = _get_cpu_list()
+    if platform.system() == "Darwin":
+        return cpu_list
+
+    global_allowed_cpu_id_list = os.sched_getaffinity(0)
+    logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list]
+
+    return logical_cpu_list
+
+
+def get_visible_memory_node() -> list[int]:
+    if platform.system() == "Darwin":
+        return [0]
+
+    allowed_memory_node_list = get_memory_affinity()
+
+    env_key = DEVICE_CONTROL_ENV_VAR
+    if (
+        ("VLLM_CPU_SIM_MULTI_NUMA" not in os.environ)
+        and env_key in os.environ
+        and os.environ[env_key] != ""
+    ):
+        visible_nodes = [int(s) for s in os.environ[env_key].split(",")]
+        visible_nodes = [
+            node for node in visible_nodes if node in allowed_memory_node_list
+        ]
+        return visible_nodes
+
+    return allowed_memory_node_list
+
+
+@cache
+def _get_cpu_list() -> list[LogicalCPUInfo]:
+    if platform.system() == "Darwin":
+        # For MacOS, no user-level CPU affinity and SMT, return all CPUs
+        cpu_count = os.cpu_count()
+        assert cpu_count
+        return [LogicalCPUInfo(i, i, 0) for i in range(cpu_count)]
+
+    lscpu_output = subprocess.check_output(
+        "lscpu -J -e=CPU,CORE,NODE", shell=True, text=True
+    )
+
+    # For platform without NUMA, replace '-' to '0'
+    lscpu_output = re.sub(r'"node":\s*-\s*(,|\n)', r'"node": 0\1', lscpu_output)
+
+    logical_cpu_list: list[LogicalCPUInfo] = json.loads(
+        lscpu_output, object_hook=LogicalCPUInfo.json_decoder
+    )["cpus"]
+
+    # Filter CPUs with invalid attributes
+    logical_cpu_list = [
+        x for x in logical_cpu_list if -1 not in (x.id, x.physical_core, x.numa_node)
+    ]
+
+    return logical_cpu_list

vllm/utils/ompmultiprocessing.py

@@ -5,196 +5,280 @@ Copyright (c) 2026 Red Hat Inc
 Copyright (c) 2026 Cambridge Greys Ltd
 """
 
-import json
 import os
-import platform
-import subprocess
+from collections.abc import Callable
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+import vllm.utils.cpu_resource_utils as cr_utils
+from vllm import envs
+from vllm.logger import init_logger
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.utils.cpu_resource_utils import LogicalCPUInfo
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+logger = init_logger(__name__)
 
 
-def _int(arg):
-    """Relaxed parsing of ints which handles a - instead of a number.
-    The lscpu json may contain that for nodes in some cases. If that
-    is the case we parse it to zero
-    """
-    try:
-        if int(arg) >= 0:
-            return int(arg)
-    except ValueError:
-        pass
-    return 0
-
-
-def parse_mask(mask):
-    """Expand a X-Y,Z list"""
-    result = []
-    for token in mask.split(","):
-        try:
-            start, finish = token.split("-")
-            if int(start) > int(finish):
-                raise IndexError("Invalid Indexes for cpu ranges")
-            for cpu in range(int(start), int(finish) + 1):
-                result.append(cpu)
-        except ValueError:
-            result.append(int(token))
-    return set(result)
-
-
-def _get_default_affinity() -> set[int]:
-    """Get the set of CPUs the process is allowed to run on."""
-    if hasattr(os, "sched_getaffinity"):
-        return os.sched_getaffinity(0)
-    # macOS does not support sched_getaffinity; fall back to cpu_count
-    cpu_count = os.cpu_count() or 1
-    return set(range(cpu_count))
-
-
-def _get_cpu_topology_json() -> bytes:
-    """Get CPU topology as JSON.
-
-    On Linux this uses ``lscpu -Je``.  On other platforms (e.g. macOS) we
-    synthesize a simple topology where every logical CPU is its own core
-    on NUMA node 0, which is sufficient for the OMP place-list builder.
-    """
-    if platform.system() == "Linux":
-        return subprocess.run(["lscpu", "-Je"], check=True, capture_output=True).stdout
-
-    # Fallback for non-Linux (macOS, etc.)
-    cpu_count = os.cpu_count() or 1
-    cpus = []
-    for i in range(cpu_count):
-        cpus.append({"cpu": str(i), "core": str(i), "node": "0"})
-    return json.dumps({"cpus": cpus}).encode()
-
-
-def enumerate_resources(resource_map, mask=None, allowed=None):
-    """Enumerate system resources"""
-    if allowed is None:
-        allowed = _get_default_affinity()
-    if mask is not None:
-        allowed = allowed & mask
-
-    try:
-        allowed_nodes = parse_mask(os.environ["CPU_VISIBLE_MEMORY_NODES"])
-    except KeyError:
-        allowed_nodes = None
-
-    lscpu: dict[str, dict] = {"cpus": {}, "cores": {}, "nodes": {}}
-    for cpu in resource_map["cpus"]:
-        cpunum = int(cpu["cpu"])
-        if (
-            cpunum in allowed
-            and cpunum >= 0
-            and (allowed_nodes is None or _int(cpu["node"]) in allowed_nodes)
-        ):
-            lscpu["cpus"][cpunum] = [cpu]
-            core = _int(cpu["core"])
-            if lscpu["cores"].get(core, None) is None:
-                lscpu["cores"][core] = [cpu]
-            else:
-                lscpu["cores"][core].append(cpu)
-            node = _int(cpu["node"])
-            if lscpu["nodes"].get(node, None) is None:
-                lscpu["nodes"][node] = [cpu]
-            else:
-                lscpu["nodes"][node].append(cpu)
-    return lscpu
-
-
-def produce_cpu_list(cpus, smt=1):
-    """Produce a CPU list with/without SMT pairs - main cpu list case"""
-    mask: list[int] = []
-    for key, value in cpus.items():
-        exists = 0
-        for cpu in mask:
-            if cpu == value[0]["core"]:
-                exists += 1
-                break
-        if exists < smt:
-            mask.append(int(key))
-    return {"mask": set(mask), "available": True}
-
-
-def produce_cpu_sublist(scpus, smt=1):
-    """Produce a CPU list with/without SMT pairs - resource leaf case"""
-    cpu_list: list[dict] = []
-    for value in scpus:
-        exists = 0
-        for cpu in cpu_list:
-            if int(cpu["core"]) == int(value["core"]):
-                exists += 1
-                break
-        if exists < smt:
-            cpu_list.append(value)
-    mask = []
-    for cpu in cpu_list:
-        mask.append(int(cpu["cpu"]))
-
-    return {"mask": set(mask), "available": True}
-
-
-def create_omp_places(resources, strategy, smt=True):
-    """Parse CPU topology and generate possible CPU masks"""
-    omp_places = []
-    if strategy == "all":
-        omp_places.append(produce_cpu_list(resources["cpus"], smt))
-    elif strategy == "cores":
-        for value in resources["cores"].values():
-            omp_places.append(produce_cpu_sublist(value, smt))
-    elif strategy == "nodes":
-        for value in resources["nodes"].values():
-            omp_places.append(produce_cpu_sublist(value, smt))
-    else:
-        raise NotImplementedError("Unknown strategy")
-
-    return omp_places
-
-
-# pylint: disable=too-few-public-methods
 class OMPProcessManager:
-    """OMP aware wrapper to run mp Process()"""
+    def __init__(self, config: "VllmConfig"):
+        if not current_platform.is_cpu():
+            return
 
-    def __init__(self, strategy="nodes", smt=1, mock=None, affinity=None):
-        self.strategy = strategy
-        self.smt = smt
-        self.omp_places = []
-        vllm_mask = os.environ.get("VLLM_CPU_OMP_THREADS_BIND", None)
-        self.setup_omp = vllm_mask != "nobind"
-        if self.setup_omp:
-            omp_places = []
-            if vllm_mask is not None:
-                masks = []
-                for spec in vllm_mask.split("|"):
-                    masks.append(parse_mask(spec))
+        self.local_world_size = config.parallel_config.local_world_size
+        self.local_dp_rank = config.parallel_config.data_parallel_rank_local
+        # This is a bit tricky because the internal DP size
+        # is always 1 for non-MoE models
+        self.internal_dp_size = config.parallel_config._api_process_count
+
+        self.simulate_multi_node = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0"
+        ld_preload_str = os.getenv("LD_PRELOAD", "")
+        self.use_iomp = "libiomp" in ld_preload_str or "libomp" in ld_preload_str
+        self.use_gomp = "libgomp" in ld_preload_str
+
+        assert not (self.use_iomp and self.use_gomp)
+
+        # at least reserve 1/local_world_size(for ARM) core for scheduler
+        # proc as always use MP executor
+        # TODO: make scheduler proc sleep when idle
+        self.reserve_cpu_num = (
+            self.local_world_size
+            if current_platform.get_cpu_architecture() == CpuArchEnum.ARM
+            else 1
+        )
+        # reserve at one more core for nixl_connector under p/d case
+        if config.kv_transfer_config:
+            self.reserve_cpu_num += 1
+
+        if envs.VLLM_CPU_NUM_OF_RESERVED_CPU is not None:
+            if self.reserve_cpu_num > envs.VLLM_CPU_NUM_OF_RESERVED_CPU:
+                msg = (
+                    f"VLLM_CPU_NUM_OF_RESERVED_CPU is less than "
+                    "the minimum requirement"
+                    f": {self.reserve_cpu_num} cores"
+                )
+                logger.warning(msg=msg)
+            self.reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU
+
+        self._parse_omp_threads_bind_env()
+
+        assert not self.simulate_multi_node or self.auto_setup
+
+    @contextmanager
+    def configure_omp_envs(self, rank: int, local_rank: int):
+        if not current_platform.is_cpu() or self.skip_setup:
+            yield
+            return
+
+        envs_dict = {}
+        cpu_list = [str(i) for i in self.cpu_lists[local_rank]]
+        envs_dict["OMP_NUM_THREADS"] = str(len(cpu_list))
+        if self.use_iomp:
+            # set IOMP envs
+            cpu_list_str = ",".join(cpu_list)
+            envs_dict["KMP_AFFINITY"] = (
+                f"granularity=fine,explicit,proclist=[{cpu_list_str}]"
+            )
+            # The time(milliseconds) that a thread should wait after
+            # completing the execution of a parallel region, before sleeping.
+            envs_dict["KMP_BLOCKTIME"] = "1"
+            # Prevents the CPU to run into low performance state
+            envs_dict["KMP_TPAUSE"] = "0"
+            # Provides fine granularity parallelism
+            envs_dict["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist"
+            envs_dict["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist"
+            envs_dict["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist"
+        elif self.use_gomp:
+            # set GOMP envs
+            # likes '0 1 2 ...'
+            cpu_list_str = " ".join(cpu_list)
+            envs_dict["GOMP_CPU_AFFINITY"] = cpu_list_str
+        else:
+            # set OMP envs
+            # likes '{0,1,2,...}'
+            cpu_list_str = ",".join(cpu_list)
+            envs_dict["OMP_PLACES"] = f"{{{cpu_list_str}}}"
+            envs_dict["OMP_PROC_BIND"] = "true"
+
+        # backup envs
+        old_envs_dict = {}
+        for k in envs_dict:
+            old_envs_dict[k] = os.environ.get(k)
+
+        try:
+            # set envs
+            for k, v in envs_dict.items():
+                os.environ[k] = v
+            yield
+        finally:
+            # restore old envs
+            for k, v in old_envs_dict.items():  # type: ignore
+                if v is None:
+                    os.environ.pop(k, None)
+                else:
+                    os.environ[k] = v
+
+    def _parse_omp_threads_bind_env(self):
+        vllm_mask = envs.VLLM_CPU_OMP_THREADS_BIND
+        self.skip_setup = vllm_mask == "nobind"
+        self.auto_setup = vllm_mask == "auto"
+        self.reserved_cpu_list = []
+        self.cpu_lists = []
+
+        if self.auto_setup:
+            # auto generate CPU lists
+            cpu_arch = current_platform.get_cpu_architecture()
+            if cpu_arch == CpuArchEnum.POWERPC:
+                # For POWERPC SMT-8/4/2
+                cpu_list, reserve_list = self._get_autobind_cpu_ids(
+                    lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4]
+                )
+            elif cpu_arch in (CpuArchEnum.X86, CpuArchEnum.S390X):
+                # For x86/S390X SMT-2, use 1 logical CPU per physical core
+                cpu_list, reserve_list = self._get_autobind_cpu_ids(
+                    lambda cpus: cpus[-1:]
+                )
+            elif cpu_arch == CpuArchEnum.ARM:
+                # For AArch64, no SMT, use all logical CPU
+                cpu_list, reserve_list = self._get_autobind_cpu_ids(lambda cpus: cpus)
             else:
-                masks = [None]
-            if mock is None:
-                data = _get_cpu_topology_json()
-            else:
-                with open(mock, mode="rb") as jf:
-                    data = jf.read()
-            lscpu = json.loads(data)
-            for mask in masks:
-                resources = enumerate_resources(lscpu, mask, affinity)
-                omp_places.extend(create_omp_places(resources, strategy, smt))
-            self.omp_places = sorted(
-                omp_places,
-                key=lambda p: "{:04d}-{:04d}".format(len(p["mask"]), max(p["mask"])),
-                reverse=True,
+                cpu_list, reserve_list = [], []
+                raise RuntimeError(f"{cpu_arch} doesn't support auto CPU binding.")
+
+            for item in cpu_list:
+                self.cpu_lists.append([x.id for x in item])
+            self.reserved_cpu_list = [x.id for x in reserve_list]
+        elif not self.skip_setup:
+            # user defined CPU lists
+            omp_cpuids_list = vllm_mask.split("|")
+            if self.local_dp_rank is not None:
+                local_dp_rank = self.local_dp_rank
+                world_size = self.local_world_size
+                # Rank mapping [DP, PP, TP]
+                omp_cpuids_list = omp_cpuids_list[
+                    local_dp_rank * world_size : (local_dp_rank + 1) * world_size
+                ]
+
+            assert len(omp_cpuids_list) == self.local_world_size, (
+                "Given "
+                f"number of CPU id list {omp_cpuids_list} doesn't match "
+                f"local world size {self.local_world_size}."
             )
 
-    def run(self, what, *args, **kwargs):
-        """Run arg with correct OMP environment"""
-        if self.setup_omp:
-            for place in self.omp_places:
-                if place["available"]:
-                    reserve = int(os.environ.get("VLLM_CPU_NUM_OF_RESERVED_CPU", 0))
-                    place["available"] = False
-                    # pylint: disable=consider-using-f-string
-                    os.environ["OMP_PLACES"] = "{}".format(place["mask"])
-                    os.environ["OMP_NUM_THREADS"] = "{}".format(
-                        len(place["mask"]) - reserve
-                    )
-                    os.environ["OMP_PROC_BIND"] = "TRUE"
-                    return what(*args, **kwargs)
-            raise IndexError("Out of OMP places")
-        return what(*args, **kwargs)
+            # parse CPU list strings like "5,2-4" to [5, 2, 3, 4]
+            self.cpu_lists = [cr_utils.parse_id_list(s) for s in omp_cpuids_list]
+        else:
+            # skip
+            self.cpu_lists = []
+
+        msg = "OpenMP thread binding info: \n"
+        for i in range(self.local_world_size):
+            msg += f"\tlocal_rank={i}, core ids={self.cpu_lists[i]}\n"
+        msg += f"\treserved_cpus={self.reserved_cpu_list}"
+        logger.info(msg)
+
+    def _get_autobind_cpu_ids(
+        self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]
+    ) -> tuple[list[list[LogicalCPUInfo]], list[LogicalCPUInfo]]:
+        """
+        Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for
+        other processes.
+        Currently for rank N, only CPU ids on the N-th node in available NUMA
+        node list will be selected.
+        Args:
+            cpu_selector: a callable object to select CPUs from a CPU list
+            of a physical core. The input is a LogicalCPUInfo list contains
+            logical CPUs of a physical CPU, sorted by the LogicalCPUInfo.id.
+            A selected LogicalCPUInfo list should be returned.
+        """
+
+        # this memory node list has been sliced for DP offset
+        allowed_numa_nodes = cr_utils.get_visible_memory_node()
+        logical_cpu_list = cr_utils.get_allowed_cpu_list()
+
+        local_world_size = self.local_world_size
+        assert (
+            len(allowed_numa_nodes) >= local_world_size or self.simulate_multi_node
+        ), (
+            f"Not enough allowed NUMA nodes to bind threads of "
+            f"{local_world_size} local CPUWorkers. "
+            f"Allowed NUMA nodes are {allowed_numa_nodes}. "
+            "Please try to bind threads manually or decrease DP/TP/PP."
+        )
+
+        # Generate OMP CPU list for each rank
+        cpu_lists_of_ranks = []
+        reserved_cpu_list = []
+        total_cpu_num = 0
+        for local_rank in range(self.local_world_size):
+            if not self.simulate_multi_node:
+                selected_numa_node = allowed_numa_nodes[local_rank]
+                selected_logical_cpu_list = [
+                    x for x in logical_cpu_list if x.numa_node == selected_numa_node
+                ]
+            else:
+                world_size_across_dp = self.local_world_size * self.internal_dp_size
+                assert len(logical_cpu_list) >= world_size_across_dp
+                selected_logical_cpu_list = sorted(
+                    logical_cpu_list, key=lambda x: x.numa_node
+                )
+                sim_cpu_num_per_node = (
+                    len(selected_logical_cpu_list) // world_size_across_dp
+                )
+                assert self.local_dp_rank is not None
+                start_idx = (
+                    local_rank + self.local_world_size * self.local_dp_rank
+                ) * sim_cpu_num_per_node
+                selected_logical_cpu_list = selected_logical_cpu_list[
+                    start_idx : (start_idx + sim_cpu_num_per_node)
+                ]
+
+            # Select logical CPUs on same physical cores via cpu_selector
+            core_to_cpus: dict[int, list[LogicalCPUInfo]] = {}
+            for cpu_info in selected_logical_cpu_list:
+                if cpu_info.physical_core not in core_to_cpus:
+                    core_to_cpus[cpu_info.physical_core] = []
+                core_to_cpus[cpu_info.physical_core].append(cpu_info)
+            selected_logical_cpu_list = []
+            for cpu_list in core_to_cpus.values():
+                cpu_list = sorted(cpu_list, key=lambda x: x.id)
+                selected_logical_cpu_list.extend(cpu_selector(cpu_list))
+
+            # sort selected cores based on core id
+            selected_logical_cpu_list = sorted(
+                selected_logical_cpu_list, key=lambda x: x.id
+            )
+
+            cpu_lists_of_ranks.append(selected_logical_cpu_list)
+            total_cpu_num += len(selected_logical_cpu_list)
+
+        # Reserve CPUs for other processes
+        if total_cpu_num <= self.reserve_cpu_num:
+            logger.warning(
+                "Selected CPU core number (%s) "
+                "should be greater than reserved CPU core "
+                "number (%s).",
+                total_cpu_num,
+                self.reserve_cpu_num,
+            )
+            return cpu_lists_of_ranks, []
+
+        reserve_num_per_rank = [
+            self.reserve_cpu_num // self.local_world_size
+        ] * self.local_world_size
+        # last rank first
+        for i in range(
+            self.local_world_size - 1,
+            self.local_world_size - 1 - self.reserve_cpu_num % self.local_world_size,
+            -1,
+        ):
+            reserve_num_per_rank[i] += 1
+        for i in range(self.local_world_size):
+            num = reserve_num_per_rank[i]
+            if num > 0:
+                reserved_cpu_list.extend(cpu_lists_of_ranks[i][-num:])
+                cpu_lists_of_ranks[i] = cpu_lists_of_ranks[i][:-num]
+
+        return cpu_lists_of_ranks, reserved_cpu_list

vllm/utils/torch_utils.py

@@ -46,6 +46,7 @@ STR_DTYPE_TO_TORCH_DTYPE = {
     "turboquant_4bit_nc": torch.uint8,
     "turboquant_k3v4_nc": torch.uint8,
     "turboquant_3bit_nc": torch.uint8,
+    "nvfp4": torch.uint8,
 }
 
 TORCH_DTYPE_TO_NUMPY_DTYPE = {
@@ -59,17 +60,19 @@ TORCH_DTYPE_TO_NUMPY_DTYPE = {
 
 
 MODELOPT_TO_VLLM_KV_CACHE_DTYPE_MAP = {
-    # TODO: Add more modelopt kv cache dtype
-    # mappings here when it supported by some attention backend
-    # (for example supports nvfp4).
     "fp8": "fp8_e4m3",
+    "nvfp4": "nvfp4",
 }
 
 T = TypeVar("T")
 
 
 def is_quantized_kv_cache(kv_cache_dtype: str) -> bool:
-    return kv_cache_dtype.startswith("fp8") or kv_cache_dtype.endswith("per_token_head")
+    return (
+        kv_cache_dtype.startswith("fp8")
+        or kv_cache_dtype.endswith("per_token_head")
+        or kv_cache_dtype == "nvfp4"
+    )
 
 
 def kv_cache_uses_per_token_head_scales(kv_cache_dtype: str) -> bool:
@@ -299,6 +302,8 @@ def get_kv_cache_quant_algo_string(quant_cfg: dict[str, Any]) -> str | None:
                 and kv_algo.get("type") == "float"
             ):
                 kv_algo = "fp8"
+            elif kv_algo.get("num_bits") == 4 and kv_algo.get("type") == "float":
+                kv_algo = "nvfp4"
             else:
                 # Unknown/unsupported format - return "auto" as safe fallback
                 logger.warning(
@@ -375,6 +380,95 @@ def set_random_seed(seed: int | None) -> None:
         current_platform.manual_seed_all(seed)
 
 
+def nvfp4_kv_cache_full_dim(head_size: int) -> int:
+    """Packed last dim for NVFP4 KV cache: fp4 data + fp8 block scales."""
+    return head_size // 2 + head_size // 16
+
+
+def _nvfp4_split_data_scale(
+    kv_side: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split a single NVFP4 KV-side buffer into data and scale views.
+
+    The input is a 4D tensor for one KV side (K or V) whose last
+    dimension is ``full_dim = data_dim + scale_dim``.  The physical
+    layout within each side is [data | scale], both packed contiguously.
+
+    Args:
+        kv_side: 4D uint8 tensor with shape
+            ``(num_pages, dim_1, dim_2, full_dim)``.
+            May be in any permutation order (NHD or HND).
+
+    Returns:
+        ``(data, scale)`` where
+        ``data`` is a uint8 view with shape
+        ``(num_pages, dim_1, dim_2, data_dim)``.
+        ``scale`` is a float8_e4m3fn view with shape
+        ``(num_pages, dim_1, dim_2, scale_dim)``.
+    """
+    num_pages = kv_side.shape[0]
+    dim_1, dim_2 = kv_side.shape[1], kv_side.shape[2]
+    full_dim = kv_side.shape[3]
+    data_dim = full_dim * 8 // 9
+    scale_dim = full_dim - data_dim
+
+    data_per_kv = dim_1 * dim_2 * data_dim
+    page_bytes = kv_side.stride(0)
+
+    # Derive inner strides from the kv_side strides, scaling by the
+    # ratio of the target dim to full_dim.  This preserves the physical
+    # layout (NHD vs HND) encoded in the input tensor's strides.
+    s1 = kv_side.stride(1) * data_dim // full_dim
+    s2 = kv_side.stride(2) * data_dim // full_dim
+    data_shape = (num_pages, dim_1, dim_2, data_dim)
+    data_strides = (page_bytes, s1, s2, 1)
+
+    s1_s = kv_side.stride(1) * scale_dim // full_dim
+    s2_s = kv_side.stride(2) * scale_dim // full_dim
+    scale_shape = (num_pages, dim_1, dim_2, scale_dim)
+    scale_strides = (page_bytes, s1_s, s2_s, 1)
+
+    base = kv_side.storage_offset()
+    data = torch.as_strided(kv_side, data_shape, data_strides, storage_offset=base)
+    scale = torch.as_strided(
+        kv_side, scale_shape, scale_strides, storage_offset=base + data_per_kv
+    ).view(torch.float8_e4m3fn)
+
+    return data, scale
+
+
+def nvfp4_kv_cache_split_views(kv_cache: torch.Tensor) -> tuple[tuple, tuple]:
+    """Split an NVFP4 KV cache tensor into data and scale views.
+
+    Accepts either a 5D tensor ``(num_pages, 2, dim_2, dim_3, full_dim)``
+    or a 4D single-side tensor ``(num_pages, dim_2, dim_3, full_dim)``.
+
+    Per-page layout: [K_data | K_scale | V_data | V_scale].
+    Each KV side is self-contained (data followed by its scale), so the
+    5D case simply splits each side independently.
+
+    The returned views are in the same dim order as the input (NHD or
+    HND), so callers get views matching whichever order they passed in.
+
+    Args:
+        kv_cache: 5D or 4D uint8 tensor where the last dimension is
+            ``full_dim = data_dim + scale_dim = 9 * head_size / 16``.
+
+    Returns:
+        For 5D input:
+            ``(k_data, v_data), (k_scale, v_scale)``
+        For 4D input (single KV side):
+            ``(data,), (scale,)``
+    """
+    if kv_cache.dim() == 4:
+        data, scale = _nvfp4_split_data_scale(kv_cache)
+        return (data,), (scale,)
+
+    k_data, k_scale = _nvfp4_split_data_scale(kv_cache[:, 0])
+    v_data, v_scale = _nvfp4_split_data_scale(kv_cache[:, 1])
+    return (k_data, v_data), (k_scale, v_scale)
+
+
 def create_kv_caches_with_random_flash(
     num_blocks: int,
     block_size: int,
@@ -401,15 +495,31 @@ def create_kv_caches_with_random_flash(
     value_caches: list[torch.Tensor] = []
 
     for _ in range(num_layers):
-        key_value_cache = torch.empty(
-            size=kv_cache_allocation_shape, dtype=dtype, device=device
-        ).permute(*stride_order)
-        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
-            key_value_cache.uniform_(-scale, scale)
-        elif cache_dtype == "fp8":
-            _generate_random_fp8(key_value_cache, -scale, scale)
+        if cache_dtype == "nvfp4":
+            # Full page dim: fp4 data + fp8 block scales per head.
+            # Per page layout: [K_data | K_scale | V_data | V_scale]
+            # Returns [:, 0] and [:, 1] like all other dtypes.
+            full_dim = nvfp4_kv_cache_full_dim(head_size)
+            nvfp4_shape = (num_blocks, 2, block_size, num_heads, full_dim)
+            nvfp4_phys = tuple(nvfp4_shape[i] for i in stride_order)
+            inv = [stride_order.index(i) for i in range(len(stride_order))]
+            key_value_cache = torch.randint(
+                0,
+                256,
+                nvfp4_phys,
+                dtype=dtype,
+                device=device,
+            ).permute(*inv)
         else:
-            raise ValueError(f"Does not support key cache of type {cache_dtype}")
+            key_value_cache = torch.empty(
+                size=kv_cache_allocation_shape, dtype=dtype, device=device
+            ).permute(*stride_order)
+            if cache_dtype in ["auto", "half", "bfloat16", "float"]:
+                key_value_cache.uniform_(-scale, scale)
+            elif cache_dtype == "fp8":
+                _generate_random_fp8(key_value_cache, -scale, scale)
+            else:
+                raise ValueError(f"Does not support key cache of type {cache_dtype}")
         key_caches.append(key_value_cache[:, 0])
         value_caches.append(key_value_cache[:, 1])
     return key_caches, value_caches

vllm/v1/attention/backends/flashinfer.py

@@ -42,7 +42,12 @@ from vllm.utils.flashinfer import (
 )
 from vllm.utils.math_utils import cdiv
 from vllm.utils.platform_utils import is_pin_memory_available
-from vllm.utils.torch_utils import is_quantized_kv_cache, is_strictly_contiguous
+from vllm.utils.torch_utils import (
+    is_quantized_kv_cache,
+    is_strictly_contiguous,
+    nvfp4_kv_cache_full_dim,
+    nvfp4_kv_cache_split_views,
+)
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
@@ -355,6 +360,10 @@ class FlashInferBackend(AttentionBackend):
         head_size: int,
         cache_dtype_str: str = "auto",
     ) -> tuple[int, ...]:
+        if cache_dtype_str == "nvfp4":
+            # Packed layout: fp4 data + fp8 block scales in last dim
+            last_dim = nvfp4_kv_cache_full_dim(head_size)
+            return (num_blocks, 2, block_size, num_kv_heads, last_dim)
         return (num_blocks, 2, block_size, num_kv_heads, head_size)
 
     @staticmethod
@@ -608,11 +617,19 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             self.cache_dtype = self.cache_config.cache_dtype
             # Cannot use self.kv_cache_spec.dtype here because kv_cache_spec
             # storage dtype may not be the same as the op dtype (uint8 vs fp8_e4m3)
-            self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
-                self.cache_dtype
-            )
+            self.is_kvcache_nvfp4 = self.cache_dtype == "nvfp4"
+            if self.is_kvcache_nvfp4:
+                # For NVFP4, kv_cache_dtype stays as the string "nvfp4"
+                # which is passed to FlashInferImpl
+                self.kv_cache_dtype = self.cache_dtype
+                raise NotImplementedError("nvfp4 KV cache is not yet supported")
+            else:
+                self.kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
+                    self.cache_dtype
+                )
         else:
             self.cache_dtype = "auto"
+            self.is_kvcache_nvfp4 = False
             assert self.kv_cache_spec.dtype == self.model_config.dtype
             self.kv_cache_dtype = self.kv_cache_spec.dtype
 
@@ -626,7 +643,13 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             can_use_trtllm
             and not vllm_config.attention_config.disable_flashinfer_q_quantization
         ):
-            self.q_data_type = self.kv_cache_dtype
+            if self.is_kvcache_nvfp4:
+                # NVFP4 KV cache uses FP8 quantized queries
+                self.q_data_type = FlashInferBackend.get_fp8_dtype_for_flashinfer(
+                    "fp8_e4m3"
+                )
+            else:
+                self.q_data_type = self.kv_cache_dtype
         else:
             self.q_data_type = self.model_config.dtype
 
@@ -1228,6 +1251,8 @@ class FlashInferImpl(AttentionImpl):
             self.sliding_window[0] if self.sliding_window is not None else -1
         )
         self.kv_cache_dtype = kv_cache_dtype
+        self.is_kvcache_nvfp4 = kv_cache_dtype == "nvfp4"
+        self.fp4_data_dim = head_size // 2 if self.is_kvcache_nvfp4 else 0
         self.logits_soft_cap = logits_soft_cap
         self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
 
@@ -1406,7 +1431,16 @@ class FlashInferImpl(AttentionImpl):
         num_prefill_tokens = attn_metadata.num_prefill_tokens
 
         stride_order = FlashInferBackend.get_kv_cache_stride_order()
-        kv_cache_permute = kv_cache.permute(*stride_order)
+        kv_cache_permute = kv_cache.permute(*stride_order)  # HND and contiguous
+
+        # For NVFP4, the kv_cache last dim is full_dim (data + scale packed).
+        # Split into correctly-strided data and scale views.
+        nvfp4_kv_data = None
+        nvfp4_kv_block_scales = None
+        if self.is_kvcache_nvfp4:
+            nvfp4_kv_data, nvfp4_kv_block_scales = nvfp4_kv_cache_split_views(
+                kv_cache_permute
+            )
 
         use_dcp = self.dcp_world_size > 1
 
@@ -1490,8 +1524,20 @@ class FlashInferImpl(AttentionImpl):
                     assert self.o_sf_scale is None
                     out = output[num_decode_tokens:]
 
-                if attn_metadata.q_data_type != FP8_DTYPE and is_quantized_kv_cache(
-                    self.kv_cache_dtype
+                prefill_kv_block_scales = None
+                if self.is_kvcache_nvfp4:
+                    # NVFP4 trtllm-gen kernel requires FP8 query.
+                    assert attn_metadata.q_data_type == FP8_DTYPE, (
+                        "NVFP4 KV cache requires FP8 quantized queries for "
+                        "trtllm-gen prefill. Set "
+                        "disable_flashinfer_q_quantization=False."
+                    )
+                    mock_kv_cache = nvfp4_kv_data
+                    mock_block_table = block_tables_prefill
+                    prefill_kv_block_scales = nvfp4_kv_block_scales  # noqa: F841
+                elif (
+                    attn_metadata.q_data_type != FP8_DTYPE
+                    and self.kv_cache_dtype.startswith("fp8")
                 ):
                     # TRTLLM prefill attention does not support BF16 Q
                     # and fp8 kv cache. So to enable prefill attention
@@ -1636,7 +1682,9 @@ class FlashInferImpl(AttentionImpl):
 
                 trtllm_batch_decode_with_kv_cache(
                     query=decode_query,
-                    kv_cache=kv_cache_permute,
+                    kv_cache=nvfp4_kv_data
+                    if self.is_kvcache_nvfp4
+                    else kv_cache_permute,
                     workspace_buffer=workspace_buffer,
                     block_tables=block_tables_decode,
                     seq_lens=seq_lens_decode,
@@ -1667,11 +1715,13 @@ class FlashInferImpl(AttentionImpl):
             # and value[:num_actual_tokens] because the reshape_and_cache_flash
             # op uses the slot_mapping's shape to determine the number of
             # actual tokens.
+            k_cache = kv_cache[:, 0]
+            v_cache = kv_cache[:, 1]
             torch.ops._C_cache_ops.reshape_and_cache_flash(
                 key,
                 value,
-                kv_cache[:, 0],
-                kv_cache[:, 1],
+                k_cache,
+                v_cache,
                 slot_mapping,
                 self.kv_cache_dtype,
                 layer._k_scale,

vllm/v1/attention/backends/mla/cutlass_fa3_sparse.py

@@ -0,0 +1,633 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CUTLASS FA3 Sparse MLA Attention Backend for vLLM.
+
+This backend uses the vendored CUTLASS FlashAttention3 Sm90 kernel from
+sgl-attn to implement sparse MLA attention for DeepSeek-V3.2 and similar
+models on SM90 (Hopper) GPUs.
+
+Key differences from FlashMLASparseBackend:
+  - Uses BF16 KV cache (576 bytes/token) instead of FP8 (656 bytes/token)
+  - No head padding needed (FA3 handles arbitrary head counts natively)
+  - Accepts Q_rope and Q_nope (qv) separately (no ConcatMLAQ kernel)
+  - 3 sub-kernels: scheduler + main attention + combine
+  - ~4x faster per transformer block (~16us vs ~64us)
+
+All execution modes (decode, prefill, mixed) are handled identically:
+each token is treated as an independent batch element with seqlen=1.
+This simplifies metadata building and CUDA graph support.
+
+Backend priority: Highest for SM90 with kv_cache_dtype="auto".
+Graceful fallback to FlashMLA Sparse when FP8 cache requested or non-SM90.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, ClassVar
+
+import numpy as np
+import torch
+
+from vllm.v1.attention.backend import (
+    AttentionBackend,
+    AttentionCGSupport,
+    AttentionLayer,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+    CommonAttentionMetadata,
+    SparseMLAAttentionImpl,
+)
+from vllm.v1.attention.ops.cutlass_fa3 import is_cutlass_fa3_available
+
+logger = logging.getLogger(__name__)
+
+# Maximum batch size (number of tokens) for which CUTLASS FA3 is used.
+# For larger batch sizes, fall back to FlashMLA BF16 sparse prefill kernel.
+# FA3 is ~4x faster than FlashMLA for small batches (bs<=16) but regresses
+# for larger batches due to higher per-token overhead from the 3-kernel
+# launch pattern (scheduler + main + combine) and page_size=1 layout.
+MAX_BATCH_SIZE_FOR_FA3 = 16
+
+# FlashMLA sparse prefill kernel requires num_heads padded to this multiple
+# on SM90 (Hopper). SM100 (Blackwell) requires 128.
+_FLASHMLA_SM90_HEAD_PADDING = 64
+
+# Check if FlashMLA BF16 sparse kernel is available for fallback
+_flashmla_sparse_available = False
+try:
+    from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd
+
+    _flashmla_sparse_available = True
+except (ImportError, Exception):
+    pass
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.config.cache import CacheDType
+    from vllm.model_executor.layers.linear import ColumnParallelLinear
+    from vllm.platforms.interface import DeviceCapability
+    from vllm.v1.kv_cache_interface import AttentionSpec
+
+
+# ─── Backend Class ────────────────────────────────────────────────────
+
+
+class CutlassFA3MLASparseBackend(AttentionBackend):
+    """CUTLASS FA3 sparse MLA for SM90 (Hopper). BF16 KV cache only.
+
+    When FP8 cache is requested, vLLM's backend selection falls back to
+    FlashMLASparseBackend automatically since this backend only supports
+    kv_cache_dtype="auto" (which maps to BF16 for MLA).
+    """
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes() -> list[int]:
+        return [64]
+
+    @staticmethod
+    def get_name() -> str:
+        return "CUTLASS_FA3_MLA_SPARSE"
+
+    @staticmethod
+    def get_builder_cls() -> type[CutlassFA3MLASparseMetadataBuilder]:
+        return CutlassFA3MLASparseMetadataBuilder
+
+    @staticmethod
+    def get_impl_cls() -> type[CutlassFA3MLASparseImpl]:
+        return CutlassFA3MLASparseImpl
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [576]
+
+    @classmethod
+    def is_mla(cls) -> bool:
+        return True
+
+    @classmethod
+    def is_sparse(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
+        return capability.major == 9
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+        cache_dtype_str: str = "auto",
+    ) -> tuple[int, ...]:
+        # BF16 cache: 576 bf16 elements per token = 1152 bytes
+        # Layout per token: [kv_c_normed(512 bf16) | k_pe(64 bf16)]
+        return (num_blocks, block_size, head_size)
+
+    @classmethod
+    def validate_configuration(
+        cls,
+        head_size: int,
+        dtype: torch.dtype,
+        kv_cache_dtype: CacheDType | None,
+        block_size: int | None,
+        use_mla: bool,
+        has_sink: bool,
+        use_sparse: bool,
+        use_mm_prefix: bool,
+        use_per_head_quant_scales: bool,
+        device_capability: DeviceCapability,
+        attn_type: str,
+        use_non_causal: bool = False,
+    ) -> list[str]:
+        invalid = super().validate_configuration(
+            head_size=head_size,
+            dtype=dtype,
+            kv_cache_dtype=kv_cache_dtype,
+            block_size=block_size,
+            use_mla=use_mla,
+            has_sink=has_sink,
+            use_sparse=use_sparse,
+            use_mm_prefix=use_mm_prefix,
+            use_per_head_quant_scales=use_per_head_quant_scales,
+            device_capability=device_capability,
+            attn_type=attn_type,
+            use_non_causal=use_non_causal,
+        )
+        if not is_cutlass_fa3_available():
+            invalid.append("_cutlass_fa3_C not available (requires CUDA >= 12.4, SM90)")
+        return invalid
+
+
+# ─── Metadata ─────────────────────────────────────────────────────────
+
+
+@dataclass
+class CutlassFA3MLASparseMetadata(AttentionMetadata):
+    """Flat metadata for CUTLASS FA3 sparse MLA attention.
+
+    ALL tokens (decode/prefill/mixed) are treated as independent batch
+    elements with seqlen=1. There are no nested Decode/Prefill sub-objects.
+
+    This simplification is valid because:
+    - Sparse MLA always routes through forward_mqa (not forward_mha)
+    - Each token independently selects its top-K KV positions
+    - The FA3 kernel handles variable-length sequences via cu_seqlens
+    """
+
+    num_reqs: int
+    max_query_len: int
+    max_seq_len: int
+    num_actual_tokens: int
+    query_start_loc: torch.Tensor
+    slot_mapping: torch.Tensor
+    block_table: torch.Tensor  # [num_reqs, max_blocks_per_req] int32
+    req_id_per_token: torch.Tensor  # [T] int32
+
+    block_size: int = 64
+    topk_tokens: int = 2048
+
+    # FA3-specific metadata (pre-allocated for CUDA graph safety)
+    cache_seqlens: torch.Tensor | None = None  # [T] int32
+    cu_seqlens_q: torch.Tensor | None = None  # [T+1] int32
+    cu_seqlens_k: torch.Tensor | None = None  # [T+1] int32
+
+    # For MLAAttention.forward_impl() routing: sparse -> all MQA
+    # Setting num_decodes = num_reqs ensures all tokens go through
+    # the forward_mqa path (no MHA prefill path).
+    num_decodes: int | None = 0
+    num_decode_tokens: int | None = 0
+    num_prefills: int | None = 0
+    num_prefill_tokens: int | None = 0
+
+
+# ─── Metadata Builder ─────────────────────────────────────────────────
+
+
+class CutlassFA3MLASparseMetadataBuilder(
+    AttentionMetadataBuilder[CutlassFA3MLASparseMetadata]
+):
+    """Builds CutlassFA3MLASparseMetadata from CommonAttentionMetadata.
+
+    Key design choices:
+    - Pre-allocates GPU buffers in __init__ for CUDA graph compatibility
+    - All tokens (decode + prefill) treated as independent seqlen=1 elements
+    - Uses in-place .copy_() for buffer updates (safe for CUDA graph replay)
+    """
+
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+
+    def __init__(
+        self,
+        kv_cache_spec: AttentionSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.topk_tokens = 2048
+        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+        self.block_size = kv_cache_spec.block_size
+
+        # Enable speculative decoding support
+        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
+
+        # Pre-allocate GPU buffers (persist across CUDA graph replays).
+        # These are updated in-place via .copy_() before each replay.
+        self.req_id_buf = torch.zeros(max_tokens, dtype=torch.int32, device=device)
+        self.cache_seqlens_buf = torch.ones(
+            max_tokens, dtype=torch.int32, device=device
+        )
+        self.cu_seqlens_q_buf = torch.arange(
+            0, max_tokens + 1, dtype=torch.int32, device=device
+        )
+        self.cu_seqlens_k_buf = torch.zeros(
+            max_tokens + 1, dtype=torch.int32, device=device
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata,
+        fast_build: bool = False,
+    ) -> CutlassFA3MLASparseMetadata:
+        """Build metadata from common attention metadata.
+
+        Converts the request-level metadata into per-token flat metadata:
+        - req_id_per_token: maps each token to its request index
+        - cache_seqlens: min(seq_len, topk) per token (topk-clipped)
+        - cu_seqlens_q: [0, 1, 2, ..., T] (each token = seqlen 1)
+        - cu_seqlens_k: cumsum of cache_seqlens
+        """
+        cm = common_attn_metadata
+        T = cm.num_actual_tokens
+        starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
+        seg_lens = np.diff(starts)
+
+        # req_id_per_token: map each token -> request index
+        req_ids = np.repeat(np.arange(len(seg_lens), dtype=np.int32), seg_lens)
+        # CUDA graph padding fix: T = cm.num_actual_tokens may include
+        # padding tokens (e.g., T=32 when only 31 real tokens exist).
+        # The computed req_ids array has sum(seg_lens) elements which
+        # equals the real (unpadded) token count. We must:
+        # 1) Zero-fill the entire buffer first (safe default for padding)
+        # 2) Copy only the actual data using req_ids.shape[0]
+        # 3) Slice to padded T for the metadata return
+        # This matches the pattern used by FlashMLASparseMetadataBuilder,
+        # FlashInferMLASparseMetadataBuilder, and all other sparse backends.
+        actual_tokens = req_ids.shape[0]
+        self.req_id_buf.fill_(0)
+        self.req_id_buf[:actual_tokens].copy_(
+            torch.from_numpy(req_ids), non_blocking=True
+        )
+
+        # cache_seqlens: UPPER BOUND = min(seq_len, topk) per token.
+        # NOTE: This is a per-REQUEST uniform value, NOT the correct
+        # per-token causal seqlen. For prefill, token i at position p
+        # can only attend to min(p+1, topk) entries, but this gives
+        # all tokens min(seq_len, topk). The actual per-token
+        # cache_seqlens is computed in forward_mqa() using valid_counts
+        # from the index conversion kernel, which correctly reflects
+        # the number of valid KV entries per token.
+        seq_lens_np = np.asarray(cm.seq_lens_cpu, dtype=np.int32)
+        per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), self.topk_tokens)
+        # Same CUDA graph padding fix: zero-fill then copy actual data.
+        # Default to 1 (safe minimum seqlen for FA3 kernel).
+        self.cache_seqlens_buf.fill_(1)
+        self.cache_seqlens_buf[:actual_tokens].copy_(
+            torch.from_numpy(per_tok_seqlens), non_blocking=True
+        )
+
+        # cu_seqlens_q: [0, 1, 2, ..., T] — each token is seqlen=1
+        cu_q = self.cu_seqlens_q_buf[: T + 1]
+
+        # cu_seqlens_k: cumsum(cache_seqlens)
+        self.cu_seqlens_k_buf[0] = 0
+        self.cu_seqlens_k_buf[1 : T + 1].copy_(
+            torch.cumsum(self.cache_seqlens_buf[:T], dim=0)
+        )
+        cu_k = self.cu_seqlens_k_buf[: T + 1]
+
+        return CutlassFA3MLASparseMetadata(
+            num_reqs=cm.num_reqs,
+            max_query_len=cm.max_query_len,
+            max_seq_len=cm.max_seq_len,
+            num_actual_tokens=T,
+            query_start_loc=cm.query_start_loc,
+            slot_mapping=cm.slot_mapping,
+            block_table=cm.block_table_tensor,
+            req_id_per_token=self.req_id_buf[:T],
+            block_size=self.block_size,
+            topk_tokens=self.topk_tokens,
+            cache_seqlens=self.cache_seqlens_buf[:T],
+            cu_seqlens_q=cu_q,
+            cu_seqlens_k=cu_k,
+            # Route ALL tokens through MQA in forward_impl
+            num_decodes=cm.num_reqs,
+            num_decode_tokens=T,
+            num_prefills=0,
+            num_prefill_tokens=0,
+        )
+
+
+# ─── Implementation ───────────────────────────────────────────────────
+
+
+class CutlassFA3MLASparseImpl(SparseMLAAttentionImpl[CutlassFA3MLASparseMetadata]):
+    """CUTLASS FA3 sparse MLA attention implementation.
+
+    This implementation replaces the FlashMLA C sparse_attn_fwd_kernel
+    with the CUTLASS FA3 Sm90 kernel from sgl-attn, providing ~4x speedup
+    per transformer block on Hopper GPUs.
+
+    Key advantages over FlashMLASparseImpl:
+    - No head padding (FA3 handles arbitrary head counts natively)
+    - No Q concatenation kernel (FA3 accepts q_rope and qv separately)
+    - BF16 KV cache (smaller footprint, no dequantization overhead)
+    - SM90 warpgroup MMA + TMA for higher compute efficiency
+    """
+
+    supports_quant_query_input: bool = False
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        # MLA Specific Arguments
+        q_lora_rank: int | None = None,
+        kv_lora_rank: int = 512,
+        qk_nope_head_dim: int = 128,
+        qk_rope_head_dim: int = 64,
+        qk_head_dim: int = 192,
+        v_head_dim: int = 128,
+        kv_b_proj: ColumnParallelLinear | None = None,
+        indexer: object | None = None,
+        q_pad_num_heads: int | None = None,
+        **kwargs,
+    ) -> None:
+        self.num_heads = num_heads  # 16 (per GPU for TP=8)
+        self.head_size = head_size  # 576 (kv_lora_rank + qk_rope_head_dim)
+        self.scale = float(scale)  # 192**-0.5
+        self.num_kv_heads = num_kv_heads  # 1 (MQA)
+        self.kv_cache_dtype = kv_cache_dtype  # "auto" (maps to BF16)
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.softmax_scale = scale
+        self.topk_tokens = 2048
+        self.num_splits = 0  # auto; CUDA-graph safe (deterministic per bs)
+        self.logits_soft_cap = float(logits_soft_cap) if logits_soft_cap else 0.0
+
+        # The indexer provides topk_indices_buffer shared across layers
+        assert indexer is not None, (
+            "CutlassFA3MLASparseImpl requires an indexer "
+            "for sparse top-K index selection"
+        )
+        self.topk_indices_buffer = indexer.topk_indices_buffer
+
+        # DCP (Decode Context Parallelism) requires softmax LSE from the
+        # attention kernel. FA3's return_softmax_lse=True is not yet wired
+        # through this backend. When DCP is needed, fall back to FlashMLA.
+        # TODO: Wire return_softmax_lse=True through forward_mqa for DCP.
+
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: CutlassFA3MLASparseMetadata,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """FA3 sparse MLA attention with batch size gating.
+
+        For batch sizes <= MAX_BATCH_SIZE_FOR_FA3 (16), uses the fast
+        CUTLASS FA3 kernel. For larger batch sizes, falls back to the
+        FlashMLA BF16 sparse prefill kernel which handles larger batches
+        more efficiently.
+
+        All execution modes (decode/prefill/mixed) are handled identically:
+        each token is an independent batch element with seqlen=1.
+
+        Input:  q = tuple(ql_nope[T, N, 512], q_pe[T, N, 64])
+        Output: (attn_out[T, N, 512], None)
+
+        The _v_up_proj in MLAAttention.forward_impl() handles the
+        subsequent .view(-1, N, kv_lora_rank) correctly for 3D output.
+        """
+        # FA3 does not yet return LSE; DCP requires it.
+        assert self.dcp_world_size <= 1, (
+            "CutlassFA3MLASparseImpl does not support DCP (dcp_world_size > 1). "
+            "Use FlashMLA Sparse instead."
+        )
+
+        # 1) Unpack Q components
+        if isinstance(q, tuple):
+            ql_nope, q_pe = q  # [T, N, 512], [T, N, 64]
+        else:
+            ql_nope = q[..., : self.kv_lora_rank]  # [T, N, 512]
+            q_pe = q[..., self.kv_lora_rank :]  # [T, N, 64]
+        T = ql_nope.shape[0]
+
+        # 2) Convert topk_indices -> global cache slot indices
+        #    Reuses vLLM's existing Triton kernel (no changes needed)
+        from vllm.v1.attention.backends.mla.sparse_utils import (
+            triton_convert_req_index_to_global_index,
+        )
+
+        global_idx, valid_counts = triton_convert_req_index_to_global_index(
+            attn_metadata.req_id_per_token,  # [T] int32
+            attn_metadata.block_table,  # [R, max_blocks] int32
+            self.topk_indices_buffer[:T],  # [T, 2048] int32
+            BLOCK_SIZE=attn_metadata.block_size,  # 64
+            NUM_TOPK_TOKENS=self.topk_tokens,  # 2048
+            return_valid_counts=True,
+        )
+        # global_idx: [T, 2048] int32 — flat cache slot IDs
+        # valid_counts: [T] int32 — number of valid (non -1) entries per token
+
+        # Replace -1 (invalid) page indices with 0 (a safe, valid page index)
+        # IN-PLACE for CUDA graph friendliness (no extra allocation).
+        global_idx.clamp_(min=0)
+
+        # Use valid_counts as cache_seqlens instead of metadata.cache_seqlens.
+        # CRITICAL FIX (Issue #1): metadata cache_seqlens = min(seq_len, topk)
+        # can exceed actual valid topk entries for prefill tokens.
+        valid_counts.clamp_(min=1)  # in-place; min=1 for seqlen_k safety
+        cache_seqlens = valid_counts
+
+        # 3) Route to FA3 or FlashMLA based on batch size
+        #    FA3 is faster for small batches (bs<=16) but regresses for
+        #    larger batches. FlashMLA BF16 sparse prefill handles larger
+        #    batches more efficiently.
+        use_fa3 = (T <= MAX_BATCH_SIZE_FOR_FA3) or not _flashmla_sparse_available
+        if use_fa3:
+            attn_out = self._forward_fa3(
+                ql_nope,
+                q_pe,
+                kv_c_and_k_pe_cache,
+                global_idx,
+                cache_seqlens,
+                attn_metadata,
+            )
+        else:
+            attn_out = self._forward_flashmla_bf16_fallback(
+                ql_nope,
+                q_pe,
+                kv_c_and_k_pe_cache,
+                global_idx,
+                cache_seqlens,
+            )
+
+        # Output: [T, N, 512] — already 3D
+        return attn_out, None
+
+    def _forward_fa3(
+        self,
+        ql_nope: torch.Tensor,  # [T, N, 512]
+        q_pe: torch.Tensor,  # [T, N, 64]
+        kv_c_and_k_pe_cache: torch.Tensor,
+        global_idx: torch.Tensor,  # [T, 2048]
+        cache_seqlens: torch.Tensor,  # [T]
+        attn_metadata: CutlassFA3MLASparseMetadata,
+    ) -> torch.Tensor:
+        """CUTLASS FA3 kernel path — fast for small batch sizes (bs<=16).
+
+        Accepts Q_rope and Q_nope (qv) separately, no head padding needed.
+        Uses page_size=1 paged KV format with split-KV parallelism.
+        """
+        T = ql_nope.shape[0]
+        S = kv_c_and_k_pe_cache.shape[0] * kv_c_and_k_pe_cache.shape[1]
+        kv_flat = kv_c_and_k_pe_cache.reshape(S, self.head_size)  # [S, 576]
+
+        # Split NoPE and RoPE, reshape for FA3 paged format (page_size=1)
+        c_kv = kv_flat[:, : self.kv_lora_rank].reshape(
+            S, 1, 1, self.kv_lora_rank
+        )  # [S, 1, 1, 512]
+        k_rope = kv_flat[:, self.kv_lora_rank :].reshape(
+            S, 1, 1, self.qk_rope_head_dim
+        )  # [S, 1, 1, 64]
+
+        from vllm.v1.attention.ops.cutlass_fa3 import flash_attn_with_kvcache
+
+        attn_out = flash_attn_with_kvcache(
+            q=q_pe,  # [T, N, 64]
+            k_cache=k_rope,  # [S, 1, 1, 64]
+            v_cache=c_kv,  # [S, 1, 1, 512]
+            qv=ql_nope,  # [T, N, 512]
+            page_table=global_idx,  # [T, 2048]
+            cache_seqlens=cache_seqlens,  # [T]
+            cu_seqlens_q=attn_metadata.cu_seqlens_q,  # [T+1]
+            cu_seqlens_k_new=None,
+            max_seqlen_q=1,
+            softmax_scale=self.softmax_scale,  # 192**-0.5
+            causal=True,
+            window_size=(-1, -1),
+            softcap=self.logits_soft_cap,
+            num_splits=self.num_splits,
+        )
+        return attn_out  # [T, N, 512]
+
+    def _forward_flashmla_bf16_fallback(
+        self,
+        ql_nope: torch.Tensor,  # [T, N, 512]
+        q_pe: torch.Tensor,  # [T, N, 64]
+        kv_c_and_k_pe_cache: torch.Tensor,
+        global_idx: torch.Tensor,  # [T, 2048]
+        cache_seqlens: torch.Tensor,  # [T]
+    ) -> torch.Tensor:
+        """FlashMLA BF16 sparse prefill fallback — for larger batch sizes.
+
+        Used when T > MAX_BATCH_SIZE_FOR_FA3 (16). The FlashMLA BF16 sparse
+        prefill kernel handles larger batches more efficiently than FA3's
+        3-kernel launch pattern (scheduler + main + combine).
+
+        This path:
+        1. Concatenates Q components: [ql_nope | q_pe] -> [T, N, 576]
+        2. Pads heads to 64 (SM90 FlashMLA requirement)
+        3. Reshapes KV cache to [S, 1, 576] (flattened, MQA format)
+        4. Reshapes indices to [T, 1, topk] (MQA format)
+        5. Calls flash_mla_sparse_fwd with topk_length for valid bounds
+        6. Unpads output heads back to N
+
+        The BF16 KV cache format [kv_c_normed(512) | k_pe(64)] is identical
+        between FA3 and FlashMLA, so no cache format conversion is needed.
+        """
+        T = ql_nope.shape[0]
+        N = self.num_heads
+
+        # 1) Concatenate Q: [ql_nope(512) | q_pe(64)] -> [T, N, 576]
+        q_concat = torch.cat([ql_nope, q_pe], dim=-1)  # [T, N, 576]
+
+        # 2) Pad heads to _FLASHMLA_SM90_HEAD_PADDING (64 for SM90)
+        padded_heads = _FLASHMLA_SM90_HEAD_PADDING
+        if padded_heads > N:
+            q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1]))
+            q_padded[:, :N, :] = q_concat
+            q_concat = q_padded
+
+        # 3) Reshape KV cache: (num_blocks, block_size, 576) -> (S, 1, 576)
+        S = kv_c_and_k_pe_cache.shape[0] * kv_c_and_k_pe_cache.shape[1]
+        kv = kv_c_and_k_pe_cache.reshape(S, 1, self.head_size)  # [S, 1, 576]
+
+        # 4) Reshape indices for MQA: (T, 2048) -> (T, 1, 2048)
+        indices = global_idx.unsqueeze(1)  # [T, 1, 2048]
+
+        # 5) Call FlashMLA BF16 sparse prefill kernel
+        #    NOTE: Unlike FlashMLASparseImpl._bf16_flash_mla_kernel which does
+        #    not pass topk_length (it relies on all indices being valid), we
+        #    pass topk_length=valid_counts because our indices have been
+        #    clamped (global_idx.clamp_(min=0)), so entries beyond valid_counts
+        #    are 0 (valid but irrelevant data). topk_length prevents the kernel
+        #    from processing these clamped entries, saving compute and ensuring
+        #    correctness.
+        output = flash_mla_sparse_fwd(
+            q_concat,  # [T, padded_heads, 576]
+            kv,  # [S, 1, 576]
+            indices,  # [T, 1, 2048]
+            self.softmax_scale,  # 192**-0.5
+            d_v=self.kv_lora_rank,  # 512
+            topk_length=cache_seqlens,  # [T] valid entry counts
+        )[0]  # extract output tensor from (output, max_logits, lse) tuple
+
+        # 6) Unpad heads: (T, padded_heads, 512) -> (T, N, 512)
+        return output[:, :N, :]
+
+    def do_kv_cache_update(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        kv_cache_dtype: str,
+        k_scale: torch.Tensor,
+    ) -> None:
+        """BF16 KV cache write using existing vLLM kernel.
+
+        kv_cache_dtype MUST be "auto" which maps to Fp8KVCacheDataType::kAuto
+        in the C++ dispatch, performing a direct BF16 copy (no quantization).
+        Passing "bfloat16" would crash because concat_and_cache_mla expects
+        the "auto" string for the non-quantized path.
+        """
+        if kv_cache.numel() == 0:
+            return
+        from vllm import _custom_ops as ops
+
+        ops.concat_and_cache_mla(
+            kv_c_normed,
+            k_pe.squeeze(1),
+            kv_cache,
+            slot_mapping.flatten(),
+            kv_cache_dtype="auto",
+            scale=k_scale,
+        )

vllm/v1/attention/backends/registry.py

@@ -73,6 +73,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     FLASHMLA_SPARSE = (
         "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend"
     )
+    CUTLASS_FA3_MLA_SPARSE = (
+        "vllm.v1.attention.backends.mla.cutlass_fa3_sparse.CutlassFA3MLASparseBackend"
+    )
     FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
     NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
     FLEX_ATTENTION = "vllm.v1.attention.backends.flex_attention.FlexAttentionBackend"

vllm/v1/attention/ops/cutlass_fa3.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Vendored CUTLASS FA3 MLA attention kernel wrapper.
+
+This module wraps the CUTLASS FlashAttention3 Sm90 kernel from sgl-attn,
+providing a Python interface compatible with vLLM's sparse MLA attention
+backend. The kernel is vendored as a self-contained C++ extension
+(_cutlass_fa3_C) and does NOT depend on sglang, sgl_kernel, or any sgl*
+modules.
+
+The FA3 kernel supports MLA (Multi-head Latent Attention) with:
+  - Separate Q_rope and QV (Q_nope) components
+  - Paged KV cache with page_size=1
+  - Variable-length sequences via cu_seqlens
+  - Split-KV parallelism with automatic split count
+  - SM90 (Hopper) CUTLASS warpgroup MMA + TMA
+
+Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
+"""
+
+import torch
+
+from vllm.platforms import current_platform
+
+_cutlass_fa3_available = False
+if current_platform.is_cuda():
+    try:
+        import vllm._cutlass_fa3_C  # noqa: F401
+
+        _cutlass_fa3_available = True
+    except ImportError:
+        pass
+
+
+def is_cutlass_fa3_available() -> bool:
+    """Check if the CUTLASS FA3 extension is available.
+
+    Requires CUDA >= 12.4 and SM90 (Hopper) GPU.
+    """
+    return _cutlass_fa3_available
+
+
+def flash_attn_with_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k: torch.Tensor | None = None,
+    v: torch.Tensor | None = None,
+    qv: torch.Tensor | None = None,
+    rotary_cos: torch.Tensor | None = None,
+    rotary_sin: torch.Tensor | None = None,
+    cache_seqlens: torch.Tensor | None = None,
+    cache_batch_idx: torch.Tensor | None = None,
+    cache_leftpad: torch.Tensor | None = None,
+    page_table: torch.Tensor | None = None,
+    cu_seqlens_q: torch.Tensor | None = None,
+    cu_seqlens_k_new: torch.Tensor | None = None,
+    max_seqlen_q: int | None = None,
+    rotary_seqlens: torch.Tensor | None = None,
+    q_descale: torch.Tensor | None = None,
+    k_descale: torch.Tensor | None = None,
+    v_descale: torch.Tensor | None = None,
+    softmax_scale: float | None = None,
+    causal: bool = False,
+    window_size: tuple[int, int] = (-1, -1),
+    attention_chunk: int | None = None,
+    softcap: float = 0.0,
+    rotary_interleaved: bool = True,
+    scheduler_metadata: torch.Tensor | None = None,
+    num_splits: int = 0,
+    pack_gqa: bool | None = None,
+    sm_margin: int = 0,
+    return_softmax_lse: bool = False,
+    sinks: torch.Tensor | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """CUTLASS FA3 attention with paged KV cache for MLA.
+
+    MLA mode shapes (DeepSeek-V3.2, varlen mode with cu_seqlens_q):
+      q:       [T, N, 64]      -- RoPE query (total_q=T, heads=N, dim=64)
+      qv:      [T, N, 512]     -- NoPE query (total_q=T, heads=N, dim_v=512)
+      k_cache: [S, 1, 1, 64]   -- Paged RoPE keys (pages, pg_sz=1, kv_h=1, d)
+      v_cache: [S, 1, 1, 512]  -- Paged NoPE latent (pages, pg_sz=1, kv_h=1, dv)
+      page_table: [T, topk]    -- Global cache slot indices per token
+
+    The FA3 kernel internally computes:
+      score = Q_rope @ K_rope^T + QV(Q_nope) @ V_cache(C_KV)^T
+      output = softmax(score * scale) @ V_cache(C_KV)
+
+    FA3 MLA constraints (from flash_api.cpp):
+      headdim_qk <= 64, headdim_v >= 256, SM90 only, BF16/FP16 only
+
+    FA3 produces 3 sub-kernels:
+      1. prepare_varlen_num_blocks_kernel (scheduler)
+      2. FlashAttnFwdSm90 (main attention, TMA+WGMMA)
+      3. FlashAttnFwdCombine (split-KV merge, when num_splits > 1)
+
+    Args:
+        q: Query tensor for RoPE component.
+        k_cache: Paged K cache (RoPE keys).
+        v_cache: Paged V cache (NoPE latent).
+        qv: Query tensor for NoPE/value component (MLA specific).
+        page_table: Page table mapping tokens to cache slots.
+        cache_seqlens: Number of valid KV entries per batch element.
+        cu_seqlens_q: Cumulative query sequence lengths.
+        cu_seqlens_k_new: Cumulative KV sequence lengths.
+        max_seqlen_q: Maximum query sequence length.
+        softmax_scale: Softmax scale factor (default: q.shape[-1]**-0.5).
+        causal: Whether to apply causal masking.
+        window_size: (left, right) attention window sizes.
+        softcap: Logits soft cap value (0.0 = disabled).
+        num_splits: Number of split-KV splits (0 = auto).
+        return_softmax_lse: Whether to return log-sum-exp values.
+
+    Returns:
+        Attention output tensor, or tuple of (output, softmax_lse).
+    """
+    assert _cutlass_fa3_available, (
+        "CUTLASS FA3 requires CUDA >= 12.4 and SM90 (Hopper) GPU. "
+        "The _cutlass_fa3_C extension was not compiled or could not be loaded."
+    )
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** (-0.5)
+    attention_chunk_val = 0 if attention_chunk is None else int(attention_chunk)
+
+    out, softmax_lse, *rest = torch.ops._cutlass_fa3_C.fwd.default(
+        q,  # 0: q
+        k_cache,  # 1: k (paged KV cache)
+        v_cache,  # 2: v (paged KV cache)
+        k,  # 3: k_new
+        v,  # 4: v_new
+        qv,  # 5: q_v (MLA NoPE query)
+        None,  # 6: out buffer
+        cu_seqlens_q,  # 7: cu_seqlens_q
+        None,  # 8: cu_seqlens_k
+        cu_seqlens_k_new,  # 9: cu_seqlens_k_new
+        None,  # 10: seqused_q
+        cache_seqlens,  # 11: seqused_k
+        max_seqlen_q,  # 12: max_seqlen_q
+        None,  # 13: max_seqlen_k
+        page_table,  # 14: page_table
+        cache_batch_idx,  # 15: kv_batch_idx
+        cache_leftpad,  # 16: leftpad_k
+        rotary_cos,  # 17: rotary_cos
+        rotary_sin,  # 18: rotary_sin
+        rotary_seqlens,  # 19: seqlens_rotary
+        q_descale,  # 20: q_descale
+        k_descale,  # 21: k_descale
+        v_descale,  # 22: v_descale
+        softmax_scale,  # 23: softmax_scale
+        causal,  # 24: is_causal
+        window_size[0],  # 25: window_size_left
+        window_size[1],  # 26: window_size_right
+        attention_chunk_val,  # 27: attention_chunk
+        softcap,  # 28: softcap
+        rotary_interleaved,  # 29: is_rotary_interleaved
+        scheduler_metadata,  # 30: scheduler_metadata
+        num_splits,  # 31: num_splits
+        pack_gqa,  # 32: pack_gqa
+        sm_margin,  # 33: sm_margin
+        sinks,  # 34: sinks
+    )
+    return (out, softmax_lse) if return_softmax_lse else out

vllm/v1/executor/multiproc_executor.py

@@ -51,6 +51,7 @@ from vllm.utils.network_utils import (
     get_loopback_ip,
     get_open_port,
 )
+from vllm.utils.ompmultiprocessing import OMPProcessManager
 from vllm.utils.system_utils import (
     _maybe_force_spawn,
     decorate_logs,
@@ -169,24 +170,14 @@ class MultiprocExecutor(Executor):
                 [] if context.get_start_method() == "fork" else None
             )
 
+            # For CPU backend only, to setup OpenMP threads affinity
+            cpu_omp_manager = OMPProcessManager(self.vllm_config)
             for local_rank in range(self.local_world_size):
                 global_rank = global_start_rank + local_rank
                 is_driver_worker = self._is_driver_worker(global_rank)
-                if current_platform.is_cpu():
-                    om = current_platform.get_omp_manager()
-                    logger.info("Configured OMP PLACES %s", str(om.omp_places))
-                    unready_worker_handle = om.run(
-                        WorkerProc.make_worker_process,
-                        vllm_config=self.vllm_config,
-                        local_rank=local_rank,
-                        rank=global_rank,
-                        distributed_init_method=distributed_init_method,
-                        input_shm_handle=scheduler_output_handle,
-                        shared_worker_lock=shared_worker_lock,
-                        is_driver_worker=is_driver_worker,
-                        inherited_fds=inherited_fds,
-                    )
-                else:
+                with cpu_omp_manager.configure_omp_envs(
+                    rank=global_rank, local_rank=local_rank
+                ):
                     unready_worker_handle = WorkerProc.make_worker_process(
                         vllm_config=self.vllm_config,
                         local_rank=local_rank,

vllm/v1/kv_cache_interface.py

@@ -17,7 +17,7 @@ from vllm.logger import init_logger
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
 from vllm.utils.math_utils import cdiv
-from vllm.utils.torch_utils import get_dtype_size
+from vllm.utils.torch_utils import get_dtype_size, nvfp4_kv_cache_full_dim
 
 logger = init_logger(__name__)
 
@@ -38,11 +38,20 @@ class KVQuantMode(IntEnum):
     FP8_PER_TENSOR = 1  # per-tensor scales (current fp8 path)
     INT8_PER_TOKEN_HEAD = 2  # per-token-head dynamic scales for int8
     FP8_PER_TOKEN_HEAD = 3  # per-token-head dynamic scales for fp8
+    NVFP4 = 4  # packed fp4 data + fp8 block scales
 
     @property
     def is_per_token_head(self) -> bool:
         """True for any per-token-head quantization mode."""
-        return self >= 2
+        return self in (
+            KVQuantMode.INT8_PER_TOKEN_HEAD,
+            KVQuantMode.FP8_PER_TOKEN_HEAD,
+        )
+
+    @property
+    def is_nvfp4(self) -> bool:
+        """True for NVFP4 packed quantization mode."""
+        return self == KVQuantMode.NVFP4
 
 
 def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
@@ -51,7 +60,9 @@ def get_kv_quant_mode(kv_cache_dtype: str) -> KVQuantMode:
         return KVQuantMode.INT8_PER_TOKEN_HEAD
     if kv_cache_dtype == "fp8_per_token_head":
         return KVQuantMode.FP8_PER_TOKEN_HEAD
-    if kv_cache_dtype.startswith("fp8"):
+    if kv_cache_dtype == "nvfp4":
+        return KVQuantMode.NVFP4
+    if isinstance(kv_cache_dtype, str) and kv_cache_dtype.startswith("fp8"):
         return KVQuantMode.FP8_PER_TENSOR
     return KVQuantMode.NONE
 
@@ -237,6 +248,19 @@ class FullAttentionSpec(AttentionSpec):
 
     @property
     def real_page_size_bytes(self) -> int:
+        if self.kv_quant_mode.is_nvfp4:
+            # Packed layout per head: fp4 data + fp8 block scales.
+            # fp4 data: head_size//2 bytes (2 fp4 values per byte)
+            # fp8 block scale: head_size//16 bytes (1 scale per 16 elements)
+            last_dim = nvfp4_kv_cache_full_dim(
+                self.head_size
+            ) + nvfp4_kv_cache_full_dim(self.head_size_v)
+            return (
+                self.block_size
+                * self.num_kv_heads
+                * last_dim
+                * get_dtype_size(self.dtype)
+            )
         return (
             self.block_size
             * self.num_kv_heads

vllm/v1/kv_offload/cpu/spec.py

@@ -26,17 +26,13 @@ class CPUOffloadingSpec(OffloadingSpec):
 
         # calculate kv_bytes_per_offloaded_block
         assert kv_cache_config is not None
-        page_sizes = {
-            kv_cache_group.kv_cache_spec.page_size_bytes
-            for kv_cache_group in kv_cache_config.kv_cache_groups
-        }
-        assert len(page_sizes) == 1
-        page_size_bytes = page_sizes.pop()
-        kv_bytes_per_block = (
-            page_size_bytes
-            * len(kv_cache_config.kv_cache_tensors)
-            * vllm_config.parallel_config.world_size
-        )
+        if kv_cache_config.num_blocks > 0:
+            total_gpu_kv_bytes = sum(t.size for t in kv_cache_config.kv_cache_tensors)
+            kv_bytes_per_block = (
+                total_gpu_kv_bytes // kv_cache_config.num_blocks
+            ) * vllm_config.parallel_config.world_size
+        else:
+            kv_bytes_per_block = 0
 
         kv_bytes_per_offloaded_block = kv_bytes_per_block * self.block_size_factor
         self.num_blocks = (

vllm/v1/worker/cpu_model_runner.py

@@ -116,21 +116,7 @@ class CPUModelRunner(GPUModelRunner):
         logger.info("Warming up model for the compilation...")
         # Only generate graph for the generic shape
         with _set_global_compilation_settings(self.vllm_config):
-            self._dummy_run(
-                min(
-                    max(16, self.max_num_reqs),
-                    self.scheduler_config.max_num_batched_tokens,
-                )
-            )
-
-        # Warm up drafter for speculative decoding
-        if self.speculative_config and (self.speculative_config.uses_draft_model()):
-            from vllm.v1.spec_decode.draft_model import DraftModelProposer
-
-            if isinstance(self.drafter, (DraftModelProposer)):
-                logger.info("Warming up drafter model...")
-                self.drafter.dummy_run(max(16, self.max_num_reqs))
-
+            self.profile_run()
         logger.info("Warming up done.")
 
     def initialize_kv_cache(

vllm/v1/worker/cpu_worker.py

@@ -1,15 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import math
 import os
 import sys
 from typing import Any
 
+import psutil
 import torch
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.platforms import CpuArchEnum, current_platform
 from vllm.profiler.wrapper import TorchProfilerWrapper
+from vllm.utils.cpu_resource_utils import (
+    get_allowed_cpu_list,
+    get_memory_node_info,
+    get_visible_memory_node,
+)
+from vllm.utils.mem_utils import format_gib
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.worker.cpu_model_runner import CPUModelRunner
 from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
@@ -27,6 +35,46 @@ class CPUWorker(Worker):
         distributed_init_method: str,
         is_driver_worker: bool = False,
     ):
+        # TODO: use numactl for process setup
+        # TODO: optimize for `interleaved` policy
+        # Bind memory node
+        allowed_memory_nodes = get_visible_memory_node()
+        allowed_cpu_list = get_allowed_cpu_list()
+        cpu_core = allowed_cpu_list[0]
+
+        # TODO: some CI hosts are not correctly set, change to assertion
+        # after fix
+        if cpu_core.numa_node not in allowed_memory_nodes:
+            logger.warning(
+                "Node %s is not in available memory nodes %s.",
+                cpu_core.numa_node,
+                allowed_memory_nodes,
+            )
+
+        torch.ops._C.init_cpu_memory_env([cpu_core.numa_node])
+
+        memory_status = get_memory_node_info(cpu_core.numa_node)
+        memory_fraction = vllm_config.cache_config.gpu_memory_utilization
+        self.requested_cpu_memory = math.ceil(
+            memory_status.total_memory * memory_fraction
+        )
+        available_memory = memory_status.available_memory
+
+        if (
+            vllm_config.cache_config.kv_cache_memory_bytes is None
+            and self.requested_cpu_memory > available_memory
+        ):
+            raise ValueError(
+                f"Available memory on node {cpu_core.numa_node} "
+                f"({format_gib(available_memory)}/"
+                f"{format_gib(memory_status.total_memory)} GiB) on startup "
+                f"is less than desired CPU memory utilization "
+                f"({vllm_config.cache_config.gpu_memory_utilization}, "
+                f"{format_gib(self.requested_cpu_memory)} GiB). "
+                "Decrease --gpu-memory-utilization"
+                f" or reduce CPU memory used by other processes."
+            )
+
         super().__init__(
             vllm_config,
             local_rank,
@@ -103,13 +151,69 @@ class CPUWorker(Worker):
         pass
 
     def determine_available_memory(self) -> int:
-        return self.cache_config.cpu_kvcache_space_bytes or 0
+        self.model_runner.warming_up_model()
+
+        allowed_cpu_list = get_allowed_cpu_list()
+        cpu_core = allowed_cpu_list[0]
+
+        memory_status = get_memory_node_info(cpu_core.numa_node)
+        available_memory = memory_status.available_memory
+        explicit_kv_cache_size = self.cache_config.kv_cache_memory_bytes
+
+        kv_cache_size = None
+        msg = None
+        if explicit_kv_cache_size is not None:
+            if explicit_kv_cache_size > available_memory:
+                raise ValueError(
+                    f"Available memory on node {cpu_core.numa_node} "
+                    f"({format_gib(available_memory)}/"
+                    f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
+                    f" allocation is less than requested memory for kv "
+                    f"({format_gib(explicit_kv_cache_size)} GiB). "
+                    "Decrease --kv-cache-memory-bytes, VLLM_CPU_KVCACHE_SPACE, "
+                    "or reduce CPU memory used by other processes."
+                )
+            kv_cache_size = explicit_kv_cache_size
+            msg = (
+                f"Explicitly set ({format_gib(kv_cache_size)}/"
+                f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
+                f"on node {cpu_core.numa_node}."
+            )
+        else:
+            consumed_memory = psutil.Process(os.getpid()).memory_info().rss
+            requested_memory_for_kv = int(self.requested_cpu_memory - consumed_memory)
+            if (
+                requested_memory_for_kv <= 0
+                or requested_memory_for_kv > available_memory
+            ):
+                raise ValueError(
+                    f"Available memory on node {cpu_core.numa_node} "
+                    f"({format_gib(available_memory)}/"
+                    f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
+                    f" allocation is less than requested memory for kv "
+                    f"({format_gib(requested_memory_for_kv)}/"
+                    f"{format_gib(self.requested_cpu_memory)} GiB). "
+                    "Reduce CPU memory used by other processes."
+                )
+            kv_cache_size = requested_memory_for_kv
+            msg = (
+                f"Auto set ({format_gib(kv_cache_size)}/"
+                f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
+                f"on node {cpu_core.numa_node}, with "
+                f"{format_gib(self.requested_cpu_memory)} GiB requested memory"
+                f" for the worker. {format_gib(consumed_memory)} GiB"
+                f" memory was consumed by non-kv usages."
+            )
+
+        logger.info(msg)
+
+        return kv_cache_size
 
     def compile_or_warm_up_model(self) -> CompilationTimes:
         # Reset the seed to ensure that the random state is not affected by
         # the model initialization and profiling.
         set_random_seed(self.model_config.seed)
-        self.model_runner.warming_up_model()
+        # Note: the model has been compiled in determine_available_memory()
         return CompilationTimes(
             language_model=self.compilation_config.compilation_time,
             encoder=self.compilation_config.encoder_compilation_time,

vllm/v1/worker/gpu/sample/gumbel.py

@@ -74,7 +74,6 @@ def gumbel_block_argmax(
     temp_ptr,
     seeds_ptr,
     pos_ptr,
-    pos_offset,
     processed_logits_ptr,
     processed_logits_stride,
     APPLY_TEMPERATURE: tl.constexpr,
@@ -99,11 +98,12 @@ def gumbel_block_argmax(
     if temp != 0.0:
         # Calculate the seed for gumbel noise.
         seed = tl.load(seeds_ptr + req_state_idx)
-        pos = tl.load(pos_ptr + token_idx) + pos_offset
+        pos = tl.load(pos_ptr + token_idx)
         gumbel_seed = tl.randint(seed, pos)
 
-        # Use FP32 for performance.
-        u = tl.rand(gumbel_seed, block)
+        # tl.rand returns fp32, so build a true fp64 uniform from 64 random
+        # bits before applying the double-log transform.
+        u = tl_rand64(gumbel_seed, block, includes_zero=False)
         gumbel_noise = -tl.log(-tl.log(u))
 
         # Apply gumbel noise.
@@ -126,7 +126,6 @@ def _gumbel_sample_kernel(
     expanded_idx_mapping_ptr,
     seeds_ptr,
     pos_ptr,
-    pos_offset,
     temp_ptr,
     vocab_size,
     BLOCK_SIZE: tl.constexpr,
@@ -152,7 +151,6 @@ def _gumbel_sample_kernel(
         temp_ptr,
         seeds_ptr,
         pos_ptr,
-        pos_offset,
         processed_logits_ptr,
         processed_logits_stride,
         APPLY_TEMPERATURE=APPLY_TEMPERATURE,
@@ -162,33 +160,6 @@ def _gumbel_sample_kernel(
     tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value)
 
 
-@triton.jit
-def _gumbel_reduce_kernel(
-    local_argmax_ptr,
-    local_argmax_stride,
-    local_max_ptr,
-    local_max_stride,
-    sampled_ptr,
-    sampled_stride,
-    num_blocks,
-    NUM_BLOCKS_NEXT_POW2: tl.constexpr,
-):
-    token_idx = tl.program_id(0)
-    offs = tl.arange(0, NUM_BLOCKS_NEXT_POW2)
-    mask = offs < num_blocks
-
-    values = tl.load(
-        local_max_ptr + token_idx * local_max_stride + offs,
-        mask=mask,
-        other=float("-inf"),
-    )
-    _, block_idx = tl.max(values, axis=0, return_indices=True)
-    token_id = tl.load(
-        local_argmax_ptr + token_idx * local_argmax_stride + block_idx,
-    )
-    tl.store(sampled_ptr + token_idx * sampled_stride, token_id)
-
-
 def gumbel_sample(
     logits: torch.Tensor,  # [num_tokens, vocab_size]
     expanded_idx_mapping: torch.Tensor,  # [num_tokens]
@@ -197,8 +168,6 @@ def gumbel_sample(
     pos: torch.Tensor,  # [num_tokens]
     apply_temperature: bool,
     processed_logits_out: torch.Tensor | None = None,  # [num_reqs, vocab_size]
-    out: torch.Tensor | None = None,  # [num_tokens], int64
-    pos_offset: int = 0,
 ) -> torch.Tensor:
     num_tokens, vocab_size = logits.shape
     BLOCK_SIZE = 1024
@@ -217,24 +186,12 @@ def gumbel_sample(
         expanded_idx_mapping,
         seed,
         pos,
-        pos_offset,
         temperature,
         vocab_size,
         BLOCK_SIZE=BLOCK_SIZE,
         APPLY_TEMPERATURE=apply_temperature,
     )
     # NOTE(woosuk): Use int64 for later indexing.
-    if out is None:
-        out = torch.empty(num_tokens, dtype=torch.int64, device=logits.device)
-    _gumbel_reduce_kernel[(num_tokens,)](
-        local_argmax,
-        local_argmax.stride(0),
-        local_max,
-        local_max.stride(0),
-        out,
-        out.stride(0),
-        num_blocks,
-        NUM_BLOCKS_NEXT_POW2=triton.next_power_of_2(num_blocks),
-        num_warps=1,
-    )
-    return out
+    max_block_idx = local_max.argmax(dim=-1, keepdim=True)
+    sampled = local_argmax.gather(dim=-1, index=max_block_idx).view(-1)
+    return sampled

vllm/v1/worker/gpu/sample/sampler.py

@@ -39,12 +39,6 @@ class Sampler:
         self.logit_bias_state = LogitBiasState(max_num_reqs, device)
         self.bad_words_state = BadWordsState(req_states)
         self.num_speculative_tokens = num_speculative_tokens
-        # Pre-allocated ones tensor for SamplerOutput.num_sampled (1 per req
-        # in the non-rejection path). Slicing returns a view so downstream
-        # reads see a stable tensor without a kernel launch per call.
-        self._num_sampled_ones = torch.ones(
-            max_num_reqs, dtype=torch.int32, device=device
-        )
 
     def add_request(
         self, req_idx: int, prompt_len: int, sampling_params: SamplingParams
@@ -68,36 +62,21 @@ class Sampler:
         expanded_idx_mapping = input_batch.expanded_idx_mapping
         idx_mapping_np = input_batch.idx_mapping_np
         cu_num_logits_np = input_batch.cu_num_logits_np
+        expanded_local_pos = input_batch.expanded_local_pos
         pos = input_batch.positions[input_batch.logits_indices]
+        input_ids = input_batch.input_ids[input_batch.logits_indices]
 
         # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
         # that num_nans is computed before applying penalties and temperature.
         num_nans = get_num_nans(logits) if self.compute_nans else None
-
-        if self._is_sampling_params_noop(idx_mapping_np):
-            # Fast path: no per-request op modifies logits. Skip the fp32 copy,
-            # the input_ids gather (only needed by bias/penalties/bad_words),
-            # and all per-state kernel dispatches.
-            sampled = gumbel_sample(
-                logits,
-                expanded_idx_mapping,
-                self.sampling_states.temperature.gpu,
-                self.sampling_states.seeds.gpu,
-                pos,
-                apply_temperature=True,
-            )
-            processed_logits = logits
-        else:
-            input_ids = input_batch.input_ids[input_batch.logits_indices]
-            expanded_local_pos = input_batch.expanded_local_pos
-            sampled, processed_logits = self.sample(
-                logits,
-                expanded_idx_mapping,
-                idx_mapping_np,
-                pos,
-                input_ids,
-                expanded_local_pos,
-            )
+        sampled, processed_logits = self.sample(
+            logits,
+            expanded_idx_mapping,
+            idx_mapping_np,
+            pos,
+            input_ids,
+            expanded_local_pos,
+        )
 
         max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)
         if max_num_logprobs != NO_LOGPROBS:
@@ -119,7 +98,7 @@ class Sampler:
             sampled_token_ids=sampled.view(-1, 1),
             logprobs_tensors=logprobs_tensors,
             num_nans=num_nans,
-            num_sampled=self._num_sampled_ones[: input_batch.num_reqs],
+            num_sampled=input_batch.seq_lens.new_ones(input_batch.num_reqs),
         )
         return sampler_output
 
@@ -172,31 +151,6 @@ class Sampler:
             logits, expanded_idx_mapping, idx_mapping_np
         )
 
-    def _is_sampling_params_noop(self, idx_mapping_np: np.ndarray) -> bool:
-        """True iff every active request uses pure defaults (argmax / gumbel).
-
-        In that case we can bypass the bf16->fp32 copy and all the per-state
-        kernel dispatches and feed the raw logits directly into
-        gumbel_sample with APPLY_TEMPERATURE=True, which handles both
-        temperature=0 (argmax) and temperature=1 (gumbel noise) without
-        any prior in-place mutation.
-        """
-        states = self.sampling_states
-        temp_np = states.temperature.np[idx_mapping_np]
-        if not np.all((temp_np == 0.0) | (temp_np == 1.0)):
-            return False
-        if np.any(states.min_p.np[idx_mapping_np] != 0.0):
-            return False
-        if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size):
-            return False
-        if np.any(states.top_p.np[idx_mapping_np] != 1.0):
-            return False
-        if np.any(self.penalties_state.use_penalty[idx_mapping_np]):
-            return False
-        if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]):
-            return False
-        return np.all(self.bad_words_state.num_bad_words.np[idx_mapping_np] == 0)
-
     def sample(
         self,
         logits: torch.Tensor,

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

@@ -237,35 +237,23 @@ class EagleSpeculator:
             cudagraph_runtime_mode=cudagraph_runtime_mode,
             mm_inputs=mm_inputs,
         )
-        # For MTP, run_model returns the same tensor for both; the two
-        # `[last_token_indices]` gathers below would be redundant, so
-        # write once into self.hidden_states and feed compute_logits from
-        # that view. For eagle3 the two tensors differ, so we still need
-        # both gathers.
-        if last_hidden_states is hidden_states:
-            self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
-            sample_hidden_states = self.hidden_states[:num_reqs]
-        else:
-            sample_hidden_states = last_hidden_states[last_token_indices]
-            self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
+        sample_hidden_states = last_hidden_states[last_token_indices]
         logits = self.model.compute_logits(sample_hidden_states)
 
         # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
-        # used for draft and target sampling. pos_offset=1 folds the +1 into
-        # the kernel itself instead of launching a separate add kernel.
-        gumbel_sample(
+        # used for draft and target sampling.
+        self.draft_tokens[:num_reqs, 0] = gumbel_sample(
             logits,
             idx_mapping,
             self.temperature,
             self.seeds,
-            pos,
+            pos + 1,
             apply_temperature=True,
             processed_logits_out=self.draft_logits[:, 0]
             if self.draft_logits is not None
             else None,
-            out=self.draft_tokens[:num_reqs, 0],
-            pos_offset=1,
         )
+        self.hidden_states[:num_reqs] = hidden_states[last_token_indices]
         self.input_buffers.positions[:num_reqs] = pos
 
     def generate_draft(
@@ -294,24 +282,19 @@ class EagleSpeculator:
             logits = self.model.compute_logits(last_hidden_states)
 
             # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise
-            # used for draft and target sampling. pos_offset=1 folds the +1
-            # into the kernel instead of launching a separate add.
-            # Write the sampled token directly into the step-th column of
-            # draft_tokens (strided view). update_eagle_inputs below will
-            # re-read it from the same slice with a matching stride.
+            # used for draft and target sampling.
             draft_tokens = gumbel_sample(
                 logits,
                 idx_mapping,
                 self.temperature,
                 self.seeds,
-                pos,
+                pos + 1,
                 apply_temperature=True,
                 processed_logits_out=self.draft_logits[:, step]
                 if self.draft_logits is not None
                 else None,
-                out=self.draft_tokens[:num_reqs, step],
-                pos_offset=1,
             )
+            self.draft_tokens[:num_reqs, step] = draft_tokens
 
             if step < self.num_speculative_steps - 1:
                 # Update the inputs for the next step.
@@ -777,7 +760,6 @@ def _update_eagle_inputs_kernel(
     seq_lens_ptr,
     max_model_len,
     draft_tokens_ptr,
-    draft_tokens_stride,
     output_hidden_states_ptr,
     output_hidden_states_stride,
     hidden_size,
@@ -786,7 +768,7 @@ def _update_eagle_inputs_kernel(
     req_idx = tl.program_id(0)
 
     # Draft token -> Input ID.
-    draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride)
+    draft_token = tl.load(draft_tokens_ptr + req_idx)
     tl.store(input_ids_ptr + req_idx, draft_token)
 
     # Output hidden states -> Input hidden states.
@@ -831,7 +813,6 @@ def update_eagle_inputs(
         input_buffers.seq_lens,
         max_model_len,
         draft_tokens,
-        draft_tokens.stride(0),
         output_hidden_states,
         output_hidden_states.stride(0),
         hidden_size,