[Bugfix] Fix k_proj's bias for GLM-ASR (#40160)

Signed-off-by: Rishapveer Singh <singhrishapveer@gmail.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

.buildkite/test-amd.yaml

@@ -2613,6 +2613,7 @@ steps:
   - vllm/platforms/rocm.py
   commands:
   - export TORCH_NCCL_BLOCKING_WAIT=1
+  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
   - pytest -v -s tests/v1/distributed/test_dbo.py
 
 
@@ -3601,7 +3602,6 @@ steps:
   commands:
   - export TORCH_NCCL_BLOCKING_WAIT=1
   - pytest -v -s tests/distributed/test_context_parallel.py
-  - VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
   - pytest -v -s tests/v1/distributed/test_dbo.py
 
 

.buildkite/test_areas/kernels.yaml

@@ -141,6 +141,7 @@ steps:
     - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
     - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
     - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
+    - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py
     - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
     - pytest -v -s tests/kernels/moe/test_flashinfer.py
     - pytest -v -s tests/kernels/moe/test_flashinfer_moe.py

.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}")
@@ -944,11 +952,19 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
       "csrc/libtorch_stable/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu"
       "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
-      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
+      "csrc/libtorch_stable/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
+      "csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu"
+      "csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu")
     set_gencode_flags_for_srcs(
       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}")

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}")

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/libtorch_stable/ops.h

@@ -134,4 +134,27 @@ void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
                               torch::stable::Tensor& input,
                               torch::stable::Tensor& input_global_scale);
 
+void mxfp4_experts_quant(
+    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
+    torch::stable::Tensor const& input,
+    torch::stable::Tensor const& input_offset_by_experts,
+    torch::stable::Tensor const& output_scale_offset_by_experts,
+    int64_t n_experts);
+
+void silu_and_mul_mxfp4_experts_quant(
+    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
+    torch::stable::Tensor const& input,
+    torch::stable::Tensor const& input_offset_by_experts,
+    torch::stable::Tensor const& output_scale_offset_by_experts,
+    int64_t n_experts);
+
+void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
+                            const torch::stable::Tensor& a,
+                            const torch::stable::Tensor& b,
+                            const torch::stable::Tensor& a_blockscale,
+                            const torch::stable::Tensor& b_blockscales,
+                            const torch::stable::Tensor& problem_sizes,
+                            const torch::stable::Tensor& expert_offsets,
+                            const torch::stable::Tensor& sf_offsets);
+
 #endif

csrc/libtorch_stable/quantization/fp4/mxfp4_blockwise_moe_kernel.cu

@@ -0,0 +1,468 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100.
+ * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups.
+ */
+
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
+
+#include <cutlass/arch/arch.h>
+
+#include "cutlass_extensions/common.hpp"
+
+#include "cute/tensor.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/epilogue/collective/default_epilogue.hpp"
+#include "cutlass/epilogue/thread/linear_combination.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/group_array_problem_shape.hpp"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+#include "cutlass/util/packed_stride.hpp"
+#include <cassert>
+
+using namespace cute;
+
+// Offset-computation kernel for MXFP4 grouped GEMM (group size 32).
+template <typename ElementAB, typename ElementC, typename ElementSF,
+          typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
+__global__ void __mxfp4_get_group_gemm_starts(
+    ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
+    ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
+    LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int,
+    ElementAB* a_base_as_int, ElementAB* b_base_as_int,
+    ElementC* out_base_as_int, ElementSF* a_scales_base_as_int,
+    ElementSF* b_scales_base_as_int, const int32_t* expert_offsets,
+    const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
+    int64_t* a_strides, int64_t* b_strides, int64_t* c_strides,
+    const int64_t a_stride_val, const int64_t b_stride_val,
+    const int64_t c_stride_val, const int K, const int N) {
+  int64_t expert_id = threadIdx.x;
+  if (expert_id >= gridDim.x * blockDim.x) {
+    return;
+  }
+  int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
+  int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
+  int64_t group_size = 32;
+  int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
+  int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
+  int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
+  assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
+         "unexpected problem sizes");
+
+  int64_t half_k = static_cast<int64_t>(k / 2);
+  int64_t group_k = static_cast<int64_t>(k / group_size);
+  // Shape of A as uint8/byte = [M, K // 2]
+  a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
+  // Shape of B as uint8/byte = [E, N, K // 2]
+  b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
+  // Shape of C = [M, N]
+  out_offsets[expert_id] = out_base_as_int + expert_offset * n;
+  // Shape of a_scale = [sum(sf_sizes), K // group_size]
+  a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
+
+  assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
+             0 &&
+         "TMA requires 128-byte alignment");
+
+  // Shape of B scale = [E, N, K // group_size]
+  b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
+  assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
+             0 &&
+         "TMA requires 128-byte alignment");
+
+  // Initialize strides
+  a_strides[expert_id] = a_stride_val;
+  b_strides[expert_id] = b_stride_val;
+  c_strides[expert_id] = c_stride_val;
+
+  LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
+  LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
+
+  *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
+      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
+  *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
+      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
+}
+
+#define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE,               \
+                                       TENSOR_C_TYPE, C_TYPE, LayoutSFA,       \
+                                       LayoutSFB, ScaleConfig)                 \
+  else if (out_tensors.scalar_type() == TENSOR_C_TYPE) {                       \
+    __mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \
+                                  LayoutSFB, ScaleConfig>                      \
+        <<<1, num_experts, 0, stream>>>(                                       \
+            static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()),               \
+            static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()),               \
+            static_cast<C_TYPE**>(out_starts.data_ptr()),                      \
+            static_cast<SF_TYPE**>(a_scales_starts.data_ptr()),                \
+            static_cast<SF_TYPE**>(b_scales_starts.data_ptr()),                \
+            reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),               \
+            reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()),               \
+            static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()),               \
+            static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()),               \
+            static_cast<C_TYPE*>(out_tensors.data_ptr()),                      \
+            static_cast<SF_TYPE*>(a_scales.data_ptr()),                        \
+            static_cast<SF_TYPE*>(b_scales.data_ptr()),                        \
+            static_cast<int32_t*>(expert_offsets.data_ptr()),                  \
+            static_cast<int32_t*>(sf_offsets.data_ptr()),                      \
+            static_cast<int32_t*>(problem_sizes.data_ptr()),                   \
+            static_cast<int64_t*>(a_strides.data_ptr()),                       \
+            static_cast<int64_t*>(b_strides.data_ptr()),                       \
+            static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val,         \
+            b_stride_val, c_stride_val, K, N);                                 \
+  }
+
+template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
+void mxfp4_run_get_group_gemm_starts(
+    const torch::stable::Tensor& a_starts,
+    const torch::stable::Tensor& b_starts,
+    const torch::stable::Tensor& out_starts,
+    const torch::stable::Tensor& a_scales_starts,
+    const torch::stable::Tensor& b_scales_starts,
+    const torch::stable::Tensor& layout_sfa,
+    const torch::stable::Tensor& layout_sfb,
+    const torch::stable::Tensor& a_strides,
+    const torch::stable::Tensor& b_strides,
+    const torch::stable::Tensor& c_strides, int64_t a_stride_val,
+    int64_t b_stride_val, int64_t c_stride_val,
+    torch::stable::Tensor const& a_tensors,
+    torch::stable::Tensor const& b_tensors,
+    torch::stable::Tensor const& out_tensors,
+    torch::stable::Tensor const& a_scales,
+    torch::stable::Tensor const& b_scales,
+    torch::stable::Tensor const& expert_offsets,
+    torch::stable::Tensor const& sf_offsets,
+    torch::stable::Tensor const& problem_sizes, int M, int N, int K) {
+  int num_experts = (int)expert_offsets.size(0);
+  auto stream = get_current_cuda_stream(a_tensors.get_device_index());
+
+  STD_TORCH_CHECK(out_tensors.size(1) == N,
+                  "Output tensor shape doesn't match expected shape");
+  STD_TORCH_CHECK(K / 2 == b_tensors.size(2),
+                  "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
+                  " dimension must match");
+  if (false) {
+  }
+  // MXFP4 uses E8M0 (float_ue8m0_t) scale factors
+  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
+                                 torch::headeronly::ScalarType::BFloat16,
+                                 cutlass::bfloat16_t, LayoutSFA, LayoutSFB,
+                                 ScaleConfig)
+  __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t,
+                                 torch::headeronly::ScalarType::Half, half,
+                                 LayoutSFA, LayoutSFB, ScaleConfig)
+  else {
+    STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
+  }
+}
+
+template <typename OutType>
+void run_mxfp4_blockwise_scaled_group_mm_sm100(
+    torch::stable::Tensor& output, const torch::stable::Tensor& a,
+    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
+    const torch::stable::Tensor& b_blockscales,
+    const torch::stable::Tensor& problem_sizes,
+    const torch::stable::Tensor& expert_offsets,
+    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
+  using ProblemShape =
+      cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
+  using ElementType = cutlass::float_e2m1_t;
+  using ElementSFType = cutlass::float_ue8m0_t;
+  using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
+  using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>;
+
+  using ElementC = OutType;
+  using ElementD = ElementC;
+  using ElementAccumulator = float;
+  // Layout definitions
+  using LayoutA = cutlass::layout::RowMajor;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  using LayoutC = cutlass::layout::RowMajor;
+  using LayoutD = LayoutC;
+
+  static constexpr int AlignmentA = 32;
+  static constexpr int AlignmentB = 32;
+  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+  // Architecture definitions
+  using ArchTag = cutlass::arch::Sm100;
+  using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;
+  using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
+  using StageCountType = cutlass::gemm::collective::StageCountAuto;
+
+  using ClusterShape = Shape<_1, _1, _1>;
+  struct MMA1SMConfig {
+    using MmaTileShape = Shape<_128, _128, _128>;
+    using KernelSchedule =
+        cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100;
+    using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;
+  };
+
+  using CollectiveEpilogue =
+      typename cutlass::epilogue::collective::CollectiveBuilder<
+          ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
+          ClusterShape, Shape<_128, _64>, ElementAccumulator,
+          ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
+          LayoutC*, AlignmentD,
+          typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
+
+  using CollectiveMainloop =
+      typename cutlass::gemm::collective::CollectiveBuilder<
+          ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
+          ElementB, LayoutB*, AlignmentB, ElementAccumulator,
+          typename MMA1SMConfig::MmaTileShape, ClusterShape,
+          cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+              sizeof(typename CollectiveEpilogue::SharedStorage))>,
+          typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
+
+  using GemmKernel =
+      cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
+                                           CollectiveEpilogue>;
+
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
+  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
+  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
+  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
+
+  using LayoutSFA =
+      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
+  using LayoutSFB =
+      typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
+  using ScaleConfig =
+      typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
+
+  using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
+  int num_experts = static_cast<int>(expert_offsets.size(0));
+
+  torch::stable::Tensor a_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor out_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor a_scales_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_scales_ptrs =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor layout_sfa = torch::stable::empty(
+      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
+      a.device());
+  torch::stable::Tensor layout_sfb = torch::stable::empty(
+      {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt,
+      a.device());
+  torch::stable::Tensor a_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor b_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+  torch::stable::Tensor c_strides1 =
+      torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long,
+                           std::nullopt, a.device());
+
+  mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
+      a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa,
+      layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2,
+      b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale,
+      b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K);
+
+  // Create an instance of the GEMM
+  Gemm gemm_op;
+
+  UnderlyingProblemShape* problem_sizes_as_shapes =
+      static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
+
+  // Set the Scheduler info
+  cutlass::KernelHardwareInfo hw_info;
+  using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
+      PersistentTileSchedulerSm100GroupParams<
+          typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
+  typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
+  scheduler.raster_order = RasterOrderOptions::AlongM;
+  hw_info.device_id = a.get_device_index();
+  static std::unordered_map<int, int> cached_sm_counts;
+  if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
+    cached_sm_counts[hw_info.device_id] =
+        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
+            hw_info.device_id);
+  }
+  hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
+
+  // Mainloop Arguments
+  typename GemmKernel::MainloopArguments mainloop_args{
+      static_cast<const ElementType**>(a_ptrs.data_ptr()),
+      static_cast<StrideA*>(a_strides1.data_ptr()),
+      static_cast<const ElementType**>(b_ptrs.data_ptr()),
+      static_cast<StrideB*>(b_strides1.data_ptr()),
+      static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
+      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
+      static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
+      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
+
+  // Epilogue Arguments
+  typename GemmKernel::EpilogueArguments epilogue_args{
+      {},  // epilogue.thread
+      nullptr,
+      static_cast<StrideC*>(c_strides1.data_ptr()),
+      static_cast<ElementD**>(out_ptrs.data_ptr()),
+      static_cast<StrideC*>(c_strides1.data_ptr())};
+  auto& fusion_args = epilogue_args.thread;
+  // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C
+  fusion_args.alpha_ptr = nullptr;
+  fusion_args.beta_ptr = nullptr;
+  fusion_args.alpha = 1.0f;
+  fusion_args.alpha_ptr_array = nullptr;
+  fusion_args.dAlpha = {_0{}, _0{}, 0};
+  fusion_args.beta = 0.0f;
+  fusion_args.beta_ptr_array = nullptr;
+  fusion_args.dBeta = {_0{}, _0{}, 0};
+
+  // Gemm Arguments
+  typename GemmKernel::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGrouped,
+      {num_experts, problem_sizes_as_shapes, nullptr},
+      mainloop_args,
+      epilogue_args,
+      hw_info,
+      scheduler};
+
+  size_t workspace_size = Gemm::get_workspace_size(args);
+  auto workspace =
+      torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte,
+                           std::nullopt, a.device());
+  const cudaStream_t stream = get_current_cuda_stream(a.get_device_index());
+
+  auto can_implement_status = gemm_op.can_implement(args);
+  STD_TORCH_CHECK(
+      can_implement_status == cutlass::Status::kSuccess,
+      "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status);
+
+  // Run the GEMM
+  auto status = gemm_op.initialize(args, workspace.data_ptr());
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
+                  "Failed to initialize MXFP4 GEMM: status=", (int)status,
+                  " workspace_size=", workspace_size,
+                  " num_experts=", num_experts, " M=", M, " N=", N, " K=", K);
+
+  status = gemm_op.run(args, workspace.data_ptr(), stream);
+  STD_TORCH_CHECK(status == cutlass::Status::kSuccess,
+                  "Failed to run MXFP4 GEMM");
+}
+
+template <typename OutType>
+void run_mxfp4_blockwise_scaled_group_mm(
+    torch::stable::Tensor& output, const torch::stable::Tensor& a,
+    const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale,
+    const torch::stable::Tensor& b_blockscales,
+    const torch::stable::Tensor& problem_sizes,
+    const torch::stable::Tensor& expert_offsets,
+    const torch::stable::Tensor& sf_offsets, int M, int N, int K) {
+  int32_t version_num = get_sm_version_num();
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
+  if (version_num >= 100 && version_num < 120) {
+    run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>(
+        output, a, b, a_blockscale, b_blockscales, problem_sizes,
+        expert_offsets, sf_offsets, M, N, K);
+    return;
+  }
+#endif
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+      false,
+      "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ",
+      version_num, ". Required capability: 100");
+}
+
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
+constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte;
+// E8M0 scale factors stored as uint8
+constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte;
+#endif
+
+#define CHECK_TYPE(x, st, m)             \
+  STD_TORCH_CHECK(x.scalar_type() == st, \
+                  ": Inconsistency of torch::stable::Tensor type:", m)
+#define CHECK_TH_CUDA(x, m) \
+  STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
+#define CHECK_CONTIGUOUS(x, m) \
+  STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
+#define CHECK_INPUT(x, st, m) \
+  CHECK_TH_CUDA(x, m);        \
+  CHECK_CONTIGUOUS(x, m);     \
+  CHECK_TYPE(x, st, m)
+
+void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
+                            const torch::stable::Tensor& a,
+                            const torch::stable::Tensor& b,
+                            const torch::stable::Tensor& a_blockscale,
+                            const torch::stable::Tensor& b_blockscales,
+                            const torch::stable::Tensor& problem_sizes,
+                            const torch::stable::Tensor& expert_offsets,
+                            const torch::stable::Tensor& sf_offsets) {
+#if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100
+  // Input validation
+  CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a");
+  CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b");
+  // MXFP4 uses E8M0 scale factors (stored as uint8)
+  CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale");
+  CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales");
+
+  STD_TORCH_CHECK(
+      a_blockscale.dim() == 2,
+      "expected a_blockscale to be of shape [num_experts, rounded_m,"
+      " k // group_size], observed rank: ",
+      a_blockscale.dim())
+  STD_TORCH_CHECK(b_blockscales.dim() == 3,
+                  "expected b_blockscale to be of shape: "
+                  " [num_experts, n, k // group_size], observed rank: ",
+                  b_blockscales.dim())
+  STD_TORCH_CHECK(problem_sizes.dim() == 2,
+                  "problem_sizes must be a 2D tensor");
+  STD_TORCH_CHECK(problem_sizes.size(1) == 3,
+                  "problem_sizes must have the shape (num_experts, 3)");
+  STD_TORCH_CHECK(
+      problem_sizes.size(0) == expert_offsets.size(0),
+      "Number of experts in problem_sizes must match expert_offsets");
+  STD_TORCH_CHECK(
+      problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int,
+      "problem_sizes must be int32.");
+
+  int M = static_cast<int>(a.size(0));
+  int N = static_cast<int>(b.size(1));
+  int E = static_cast<int>(b.size(0));
+  int K = static_cast<int>(2 * b.size(2));
+
+  if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
+    run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
+        output, a, b, a_blockscale, b_blockscales, problem_sizes,
+        expert_offsets, sf_offsets, M, N, K);
+  } else {
+    run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>(
+        output, a, b, a_blockscale, b_blockscales, problem_sizes,
+        expert_offsets, sf_offsets, M, N, K);
+  }
+#else
+  STD_TORCH_CHECK_NOT_IMPLEMENTED(
+      false,
+      "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with "
+      "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+.");
+#endif
+}
+
+STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
+  m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm));
+}

csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu

@@ -0,0 +1,422 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+ *
+ * MXFP4 activation quantization kernel for MoE experts.
+ * Quantizes BF16/FP16 activations to MXFP4: E2M1 values with E8M0 block scales
+ * over 32-element groups.
+ *
+ * Uses PACK16 E2M1 conversion helpers (nvfp4_utils.cuh) configured for:
+ *   - Block size 32 (2 threads per SF in PACK16 mode)
+ *   - E8M0 (power-of-two) scale factors
+ *   - SF layout: [numMTiles, numKTiles, 32, 4, 4] where numKTiles=ceil(K/128)
+ */
+
+// MXFP4 requires PACK16 mode (16 elements per thread) so that
+// 2 threads cover 32-element blocks. This requires CUDA >= 12.9.
+// Must be defined before any header that (transitively) includes
+// nvfp4_utils.cuh.
+#define NVFP4_ENABLE_ELTS16 1
+
+#include <cuda.h>
+#include <cuda_runtime_api.h>
+#include <cuda_runtime.h>
+#include <cuda_fp8.h>
+
+#include <torch/csrc/stable/tensor.h>
+#include "libtorch_stable/torch_utils.h"
+#include "libtorch_stable/dispatch_utils.h"
+#include "cuda_vec_utils.cuh"
+#include "cuda_utils.h"
+
+#include "nvfp4_utils.cuh"
+static_assert(CVT_FP4_ELTS_PER_THREAD == 16,
+              "MXFP4 experts quant requires PACK16 mode (CUDA >= 12.9)");
+
+#include "launch_bounds_utils.h"
+
+namespace vllm {
+
+// MXFP4 block size constants
+static constexpr int MXFP4_SF_VEC_SIZE = 32;
+
+// For PACK16 mode (CVT_FP4_ELTS_PER_THREAD=16): 2 threads per SF
+// For PACK8 mode (CVT_FP4_ELTS_PER_THREAD=8): 4 threads per SF
+static constexpr int MXFP4_NUM_THREADS_PER_SF =
+    MXFP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD;
+
+// MXFP4 quantization kernel for experts.
+// Uses 32-element blocks with E8M0 (UE8M0) scale factors.
+// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
+// SiLU(gate)*up before quantization.
+template <class Type, bool FUSE_SILU_MUL = false,
+          bool SMALL_NUM_EXPERTS = false>
+__global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
+    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
+                          fp4_packed_t* out, uint32_t* SFout,
+                          uint32_t* input_offset_by_experts,
+                          uint32_t* output_scale_offset_by_experts,
+                          int n_experts, bool low_latency) {
+  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
+  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
+                "Vec size is not matched.");
+
+  // MXFP4: numKTiles = ceil(numCols / 128) since block_size=32, 4 SFs/tile
+  int32_t const numKTiles = (numCols + 127) / 128;
+
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
+  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
+
+  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
+       globalIdx += gridDim.x * blockDim.x) {
+    int rowIdx = globalIdx / colsPerRow;
+    int colIdx = globalIdx % colsPerRow;
+
+    int rowIdx_in_expert = 0;
+    int expert_idx = 0;
+
+    if constexpr (SMALL_NUM_EXPERTS) {
+      for (int i = 0; i < n_experts; i++) {
+        uint32_t current_offset = __ldca(&input_offset_by_experts[i]);
+        uint32_t next_offset = __ldca(&input_offset_by_experts[i + 1]);
+        if (rowIdx >= current_offset && rowIdx < next_offset) {
+          rowIdx_in_expert = rowIdx - current_offset;
+          expert_idx = i;
+          break;
+        }
+      }
+    } else {
+      uint32_t local_offsets[17];
+      for (int chunk_start = 0; chunk_start < n_experts; chunk_start += 16) {
+        *reinterpret_cast<int4*>(local_offsets) =
+            __ldca(reinterpret_cast<const int4*>(
+                &input_offset_by_experts[chunk_start]));
+        *reinterpret_cast<int4*>(local_offsets + 4) =
+            __ldca(reinterpret_cast<const int4*>(
+                &input_offset_by_experts[chunk_start + 4]));
+        *reinterpret_cast<int4*>(local_offsets + 8) =
+            __ldca(reinterpret_cast<const int4*>(
+                &input_offset_by_experts[chunk_start + 8]));
+        *reinterpret_cast<int4*>(local_offsets + 12) =
+            __ldca(reinterpret_cast<const int4*>(
+                &input_offset_by_experts[chunk_start + 12]));
+        local_offsets[16] = __ldca(&input_offset_by_experts[chunk_start + 16]);
+
+#pragma unroll
+        for (int i = 0; i < 16; i++) {
+          if (rowIdx >= local_offsets[i] && rowIdx < local_offsets[i + 1]) {
+            rowIdx_in_expert = rowIdx - local_offsets[i];
+            expert_idx = chunk_start + i;
+            break;
+          }
+        }
+      }
+    }
+
+    // Load input and optionally apply fused SiLU+Mul
+    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
+    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
+    PackedVec quant_input;
+    if constexpr (FUSE_SILU_MUL) {
+      PackedVec in_vec_up =
+          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
+      quant_input = compute_silu_mul(in_vec, in_vec_up);
+    } else {
+      quant_input = in_vec;
+    }
+
+    // In PACK16 mode, each thread outputs 16 E2M1 values = u32x2
+    int64_t outOffset = rowIdx * colsPerRow + colIdx;
+    auto& out_pos = out[outOffset];
+
+    uint32_t* SFout_in_expert =
+        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
+
+    // Use MXFP4_NUM_THREADS_PER_SF (2 for PACK16) for 32-element blocks
+    auto sf_out =
+        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
+            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
+
+    // Block E8M0 scales only; no extra tensor-level scale in this path
+    constexpr float SFScaleVal = 1.0f;
+    // UE8M0_SF=true for MXFP4 E8M0 scale factors
+    out_pos =
+        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
+            quant_input, SFScaleVal, sf_out);
+  }
+}
+
+// Large M_topk variant using shared memory for expert offsets
+template <class Type, bool FUSE_SILU_MUL = false,
+          bool SMALL_NUM_EXPERTS = false>
+__global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
+    mxfp4_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
+                          fp4_packed_t* out, uint32_t* SFout,
+                          uint32_t* input_offset_by_experts,
+                          uint32_t* output_scale_offset_by_experts,
+                          int n_experts) {
+  using PackedVec = PackedVec<Type, CVT_FP4_PACK16>;
+  static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
+                "Vec size is not matched.");
+
+  // MXFP4: numKTiles = ceil(numCols / 128)
+  int32_t const numKTiles = (numCols + 127) / 128;
+
+  extern __shared__ uint32_t shared_input_offsets[];
+
+  if constexpr (SMALL_NUM_EXPERTS) {
+    for (int i = threadIdx.x; i < n_experts + 1; i += blockDim.x) {
+      shared_input_offsets[i] = input_offset_by_experts[i];
+    }
+  } else {
+    for (int i = threadIdx.x * 4; i < n_experts; i += blockDim.x * 4) {
+      *reinterpret_cast<int4*>(&shared_input_offsets[i]) =
+          *reinterpret_cast<const int4*>(&input_offset_by_experts[i]);
+    }
+    if (threadIdx.x == 0) {
+      shared_input_offsets[n_experts] = input_offset_by_experts[n_experts];
+    }
+  }
+
+  __syncthreads();
+
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
+  int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
+
+  for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
+       globalIdx += gridDim.x * blockDim.x) {
+    int rowIdx = globalIdx / colsPerRow;
+    int colIdx = globalIdx % colsPerRow;
+
+    int rowIdx_in_expert = 0;
+    int expert_idx = 0;
+
+    // Binary search through experts using shared memory
+    int left = 0, right = n_experts - 1;
+    while (left <= right) {
+      int mid = (left + right) / 2;
+      uint32_t mid_offset = shared_input_offsets[mid];
+      uint32_t next_offset = shared_input_offsets[mid + 1];
+
+      if (rowIdx >= mid_offset && rowIdx < next_offset) {
+        rowIdx_in_expert = rowIdx - mid_offset;
+        expert_idx = mid;
+        break;
+      } else if (rowIdx < mid_offset) {
+        right = mid - 1;
+      } else {
+        left = mid + 1;
+      }
+    }
+
+    int64_t inOffset = rowIdx * inColsPerRow + colIdx;
+    PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
+    PackedVec quant_input;
+    if constexpr (FUSE_SILU_MUL) {
+      PackedVec in_vec_up =
+          reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
+      quant_input = compute_silu_mul(in_vec, in_vec_up);
+    } else {
+      quant_input = in_vec;
+    }
+
+    int64_t outOffset = rowIdx * colsPerRow + colIdx;
+    auto& out_pos = out[outOffset];
+
+    // MXFP4 has no global scale - only block-level E8M0 scale factors
+    constexpr float SFScaleVal = 1.0f;
+
+    uint32_t* SFout_in_expert =
+        SFout + output_scale_offset_by_experts[expert_idx] * numKTiles;
+
+    auto sf_out =
+        cvt_quant_to_fp4_get_sf_out_offset<uint32_t, MXFP4_NUM_THREADS_PER_SF>(
+            rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
+
+    out_pos =
+        cvt_warp_fp16_to_fp4<Type, MXFP4_NUM_THREADS_PER_SF, /*UE8M0_SF=*/true>(
+            quant_input, SFScaleVal, sf_out);
+  }
+}
+
+template <typename T, bool FUSE_SILU_MUL = false>
+void mxfp4_quant_impl(void* output, void* output_scale, void* input,
+                      void* input_offset_by_experts,
+                      void* output_scale_offset_by_experts, int m_topk, int k,
+                      int n_experts, cudaStream_t stream) {
+  int multiProcessorCount =
+      get_device_attribute(cudaDevAttrMultiProcessorCount, -1);
+
+  int const workSizePerRow = k / ELTS_PER_THREAD;
+  int const totalWorkSize = m_topk * workSizePerRow;
+  dim3 block(std::min(workSizePerRow, 512));
+  int const numBlocksPerSM =
+      vllm_runtime_blocks_per_sm(static_cast<int>(block.x));
+  dim3 grid(std::min(static_cast<int>((totalWorkSize + block.x - 1) / block.x),
+                     multiProcessorCount * numBlocksPerSM));
+  while (grid.x <= multiProcessorCount && block.x > 64) {
+    grid.x *= 2;
+    block.x = (block.x + 1) / 2;
+  }
+
+  int const blockRepeat =
+      (totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
+  if (blockRepeat > 1) {
+    size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
+    if (n_experts >= 4) {
+      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
+          <<<grid, block, shared_mem_size, stream>>>(
+              m_topk, k, reinterpret_cast<T*>(input),
+              reinterpret_cast<fp4_packed_t*>(output),
+              reinterpret_cast<uint32_t*>(output_scale),
+              reinterpret_cast<uint32_t*>(input_offset_by_experts),
+              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+              n_experts);
+    } else {
+      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true>
+          <<<grid, block, shared_mem_size, stream>>>(
+              m_topk, k, reinterpret_cast<T*>(input),
+              reinterpret_cast<fp4_packed_t*>(output),
+              reinterpret_cast<uint32_t*>(output_scale),
+              reinterpret_cast<uint32_t*>(input_offset_by_experts),
+              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+              n_experts);
+    }
+  } else {
+    if (n_experts >= 16) {
+      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false>
+          <<<grid, block, 0, stream>>>(
+              m_topk, k, reinterpret_cast<T*>(input),
+              reinterpret_cast<fp4_packed_t*>(output),
+              reinterpret_cast<uint32_t*>(output_scale),
+              reinterpret_cast<uint32_t*>(input_offset_by_experts),
+              reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+              n_experts, /* bool low_latency */ true);
+    } else {
+      mxfp4_cvt_fp16_to_fp4<T, FUSE_SILU_MUL, true><<<grid, block, 0, stream>>>(
+          m_topk, k, reinterpret_cast<T*>(input),
+          reinterpret_cast<fp4_packed_t*>(output),
+          reinterpret_cast<uint32_t*>(output_scale),
+          reinterpret_cast<uint32_t*>(input_offset_by_experts),
+          reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
+          n_experts, /* bool low_latency */ true);
+    }
+  }
+}
+
+}  // namespace vllm
+
+/*Quantization entry for mxfp4 experts quantization*/
+#define CHECK_TH_CUDA(x, m) \
+  STD_TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x, m) \
+  STD_TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
+#define CHECK_INPUT(x, m) \
+  CHECK_TH_CUDA(x, m);    \
+  CHECK_CONTIGUOUS(x, m);
+
+constexpr auto HALF = torch::headeronly::ScalarType::Half;
+constexpr auto BF16 = torch::headeronly::ScalarType::BFloat16;
+constexpr auto INT = torch::headeronly::ScalarType::Int;
+constexpr auto UINT8 = torch::headeronly::ScalarType::Byte;
+
+static constexpr int MXFP4_BLOCK_SIZE = 32;
+
+static void validate_mxfp4_experts_quant_inputs(
+    torch::stable::Tensor const& output,
+    torch::stable::Tensor const& output_scale,
+    torch::stable::Tensor const& input,
+    torch::stable::Tensor const& input_offset_by_experts,
+    torch::stable::Tensor const& output_scale_offset_by_experts,
+    int64_t n_experts, int64_t m_topk, int64_t k) {
+  CHECK_INPUT(output, "output");
+  CHECK_INPUT(output_scale, "output_scale");
+  CHECK_INPUT(input, "input");
+  CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
+  CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
+
+  STD_TORCH_CHECK(output.dim() == 2);
+  STD_TORCH_CHECK(output_scale.dim() == 2);
+  STD_TORCH_CHECK(input.dim() == 2);
+  STD_TORCH_CHECK(input_offset_by_experts.dim() == 1);
+  STD_TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
+
+  STD_TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
+  STD_TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
+  STD_TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
+  // output is uint8 (two mxfp4 values packed into one uint8)
+  // output_scale is int32 (four E8M0 values packed into one int32)
+  STD_TORCH_CHECK(output.scalar_type() == UINT8);
+  STD_TORCH_CHECK(output_scale.scalar_type() == INT);
+
+  STD_TORCH_CHECK(k % MXFP4_BLOCK_SIZE == 0, "k must be a multiple of 32");
+  STD_TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
+  STD_TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
+  STD_TORCH_CHECK(output.size(0) == m_topk);
+  STD_TORCH_CHECK(output.size(1) == k / 2);
+  int scales_k = k / MXFP4_BLOCK_SIZE;
+  // K-dimension scale columns padded to a multiple of 4 for swizzle layout
+  int padded_k = (scales_k + (4 - 1)) / 4 * 4;
+  // 4 = 4 E8M0 values packed into one int32
+  STD_TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
+}
+
+void mxfp4_experts_quant(
+    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
+    torch::stable::Tensor const& input,
+    torch::stable::Tensor const& input_offset_by_experts,
+    torch::stable::Tensor const& output_scale_offset_by_experts,
+    int64_t n_experts) {
+  auto m_topk = input.size(0);
+  auto k = input.size(1);
+
+  validate_mxfp4_experts_quant_inputs(
+      output, output_scale, input, input_offset_by_experts,
+      output_scale_offset_by_experts, n_experts, m_topk, k);
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
+
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      input.scalar_type(), "mxfp4_experts_quant_kernel", [&] {
+        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
+        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
+            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
+            input_offset_by_experts.data_ptr(),
+            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
+            stream);
+      });
+}
+
+void silu_and_mul_mxfp4_experts_quant(
+    torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
+    torch::stable::Tensor const& input,
+    torch::stable::Tensor const& input_offset_by_experts,
+    torch::stable::Tensor const& output_scale_offset_by_experts,
+    int64_t n_experts) {
+  auto m_topk = input.size(0);
+  auto k_times_2 = input.size(1);
+  STD_TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
+  auto k = k_times_2 / 2;
+
+  validate_mxfp4_experts_quant_inputs(
+      output, output_scale, input, input_offset_by_experts,
+      output_scale_offset_by_experts, n_experts, m_topk, k);
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      input.get_device_index());
+  const cudaStream_t stream = get_current_cuda_stream(input.get_device_index());
+
+  VLLM_STABLE_DISPATCH_HALF_TYPES(
+      input.scalar_type(), "silu_mul_mxfp4_experts_quant_kernel", [&] {
+        using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
+        vllm::mxfp4_quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
+            output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
+            input_offset_by_experts.data_ptr(),
+            output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
+            stream);
+      });
+}

csrc/libtorch_stable/torch_bindings.cpp

@@ -116,6 +116,12 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
       " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
 
+  // cutlass mxfp4 block scaled group GEMM (MXFP4 x MXFP4 MoE)
+  ops.def(
+      "cutlass_mxfp4_group_mm(Tensor! out, Tensor a, Tensor b,"
+      " Tensor a_blockscale, Tensor b_blockscales,"
+      " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()");
+
   // Compute NVFP4 block quantized tensor.
   ops.def(
       "scaled_fp4_quant(Tensor input,"
@@ -149,6 +155,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
       "Tensor output_scale_offset_by_experts) -> ()");
 
+  // Compute MXFP4 experts quantization (32-element blocks, E8M0 SFs).
+  ops.def(
+      "mxfp4_experts_quant(Tensor! output, Tensor! output_scale,"
+      "Tensor input, Tensor input_offset_by_experts,"
+      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
+
+  // Fused SiLU+Mul+MXFP4 experts quantization.
+  ops.def(
+      "silu_and_mul_mxfp4_experts_quant(Tensor! output, Tensor! "
+      "output_scale,"
+      "Tensor input, Tensor input_offset_by_experts,"
+      "Tensor output_scale_offset_by_experts, int n_experts) -> ()");
+
   // Fused SiLU+Mul+NVFP4 quantization.
   ops.def(
       "silu_and_mul_nvfp4_quant(Tensor! result, Tensor! result_block_scale, "
@@ -233,6 +252,9 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("silu_and_mul_scaled_fp4_experts_quant",
            TORCH_BOX(&silu_and_mul_scaled_fp4_experts_quant));
   ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
+  ops.impl("mxfp4_experts_quant", TORCH_BOX(&mxfp4_experts_quant));
+  ops.impl("silu_and_mul_mxfp4_experts_quant",
+           TORCH_BOX(&silu_and_mul_mxfp4_experts_quant));
 
   // W4A8 ops: impl registrations are in the source files
   // (w4a8_mm_entry.cu and w4a8_grouped_mm_entry.cu)

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);
+      });
+}

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

docker/Dockerfile.rocm

@@ -192,9 +192,10 @@ RUN cd /opt/rixl && mkdir -p /app/install && \
 FROM base AS build_deep
 ARG ROCSHMEM_BRANCH="ba0bf0f3"
 ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
-ARG DEEPEP_BRANCH="e84464ec"
+ARG DEEPEP_BRANCH="5d90af8b"
 ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
 ARG DEEPEP_NIC="cx7"
+ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
 ENV ROCSHMEM_DIR=/opt/rocshmem
 
 RUN git clone ${ROCSHMEM_REPO} \
@@ -202,13 +203,11 @@ RUN git clone ${ROCSHMEM_REPO} \
  && git checkout ${ROCSHMEM_BRANCH} \
  && mkdir -p projects/rocshmem/build \
  && cd projects/rocshmem/build \
- && cmake .. \
-    -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
-    -DROCM_PATH=/opt/rocm \
-    -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
-    -DUSE_EXTERNAL_MPI=OFF \
- && make -j \
- && make install
+ && bash ../scripts/build_configs/all_backends \
+      -DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
+      -DROCM_PATH=/opt/rocm \
+      -DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
+      -DUSE_EXTERNAL_MPI=OFF
 
 # Build DeepEP wheel.
 # DeepEP looks for rocshmem at ROCSHMEM_DIR.
@@ -262,30 +261,6 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
     && echo "Detected vLLM version: ${VLLM_VERSION}" \
     && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
 
-# Fail if git-based package dependencies are found in requirements files
-# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
-# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
-RUN echo "Checking for git-based packages in requirements files..." \
-    && echo "Checking common.txt for git-based packages:" \
-    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
-         echo "ERROR: Git-based packages found in common.txt:"; \
-         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
-         echo "Please publish these packages to PyPI instead of using git dependencies."; \
-         exit 1; \
-       else \
-         echo "  ✓ No git-based packages found in common.txt"; \
-       fi \
-    && echo "Checking rocm.txt for git-based packages:" \
-    && if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
-         echo "ERROR: Git-based packages found in rocm.txt:"; \
-         grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
-         echo "Please publish these packages to PyPI instead of using git dependencies."; \
-         exit 1; \
-       else \
-         echo "  ✓ No git-based packages found in rocm.txt"; \
-       fi \
-    && echo "All requirements files are clean - no git-based packages found"
-
 # Pin vLLM dependencies to exact versions of custom ROCm wheels
 # This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
 COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py

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)
 

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.
 
@@ -674,6 +682,24 @@ Speech2Text models trained specifically for Automatic Speech Recognition.
 !!! note
     `VoxtralForConditionalGeneration` requires `mistral-common[audio]` to be installed.
 
+#### Realtime Transcription
+
+Speech models that support streaming transcription via the
+[`/v1/realtime`](../serving/openai_compatible_server.md#realtime-api)
+WebSocket endpoint.
+
+| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
+| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
+| `VoxtralRealtimeGeneration` | Voxtral Realtime | `mistralai/Voxtral-Mini-4B-Realtime-2602` | | |
+| `Qwen3ASRRealtimeGeneration` | Qwen3-ASR Realtime | `Qwen/Qwen3-ASR-0.6B` | | |
+
+!!! note
+    `VoxtralRealtimeGeneration` requires `mistral-common[audio]` to be installed, and must be served with `--tokenizer-mode mistral`.
+
+    `Qwen3ASRRealtimeGeneration` is not auto-detected from `config.json`.
+    You must pass `--hf-overrides '{"architectures":["Qwen3ASRRealtimeGeneration"]}'`
+    when serving.
+
 ## Pooling Models
 
 See [this page](pooling_models/README.md) for more information on how to use pooling models.

docs/serving/openai_compatible_server.md

@@ -60,7 +60,7 @@ We currently support the following OpenAI APIs:
 - [Translation API](#translations-api) (`/v1/audio/translations`)
     - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
 - [Realtime API](#realtime-api) (`/v1/realtime`)
-    - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#transcription).
+    - Only applicable to [Automatic Speech Recognition (ASR) models](../models/supported_models.md#realtime-transcription).
 
 In addition, we have the following custom APIs:
 

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

examples/online_serving/disaggregated_serving/example_mm_serve.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Disaggregated multimodal serving: render → generate round-trip.
+
+Demonstrates the two-phase disaggregated flow:
+  1. /v1/chat/completions/render  – preprocesses a multimodal chat request
+     into token IDs and serialized tensor features.
+  2. /inference/v1/generate       – runs inference on the preprocessed tokens.
+
+The render response is passed *directly* to generate with only
+``sampling_params`` added, showing that the two endpoints compose with
+zero client-side transformation.
+
+Launch the server first:
+
+    vllm serve Qwen/Qwen3-VL-2B-Instruct \
+        --dtype bfloat16 --max-model-len 4096 --enforce-eager
+
+Then run this script:
+
+    python example_mm_serve.py
+"""
+
+import io
+
+import pybase64 as base64
+import requests
+from PIL import Image
+from transformers import AutoTokenizer
+
+BASE_URL = "http://localhost:8000"
+MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
+
+
+def make_data_url(image: Image.Image) -> str:
+    """Encode a PIL image as a base64 data URL."""
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/png;base64,{b64}"
+
+
+def main():
+    # -- Step 1: Create a test image (solid red) -------------------------
+    image = Image.new("RGB", (224, 224), color=(255, 0, 0))
+    data_url = make_data_url(image)
+    print("Created 224x224 red test image")
+
+    # -- Step 2: Render (preprocess) -------------------------------------
+    render_payload = {
+        "model": MODEL_NAME,
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image_url", "image_url": {"url": data_url}},
+                    {
+                        "type": "text",
+                        "text": "What color is this image? Answer in one word.",
+                    },
+                ],
+            }
+        ],
+    }
+
+    print("\n--- Render ---")
+    render_resp = requests.post(
+        f"{BASE_URL}/v1/chat/completions/render", json=render_payload
+    )
+    render_resp.raise_for_status()
+    render_data = render_resp.json()
+
+    print(f"Response keys: {list(render_data.keys())}")
+    print(f"Number of token_ids: {len(render_data['token_ids'])}")
+
+    features = render_data.get("features")
+    if features and features.get("kwargs_data"):
+        print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
+        for modality, items in features["kwargs_data"].items():
+            print(
+                f"  {modality}: {len(items)} item(s), "
+                f"first item type: {type(items[0])} length: {len(items[0])}"
+                if items
+                else "First item: (empty)"
+            )
+    else:
+        print("WARNING: no kwargs_data in render response")
+
+    # -- Step 3: Generate (inference) ------------------------------------
+    # Pass the render output directly — only add sampling_params.
+    generate_payload = render_data
+    generate_payload["sampling_params"] = {
+        "max_tokens": 20,
+        "temperature": 0.0,
+    }
+
+    print("\n--- Generate ---")
+    gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
+    gen_resp.raise_for_status()
+    gen_data = gen_resp.json()
+
+    # -- Step 4: Decode & print ------------------------------------------
+    output_ids = gen_data["choices"][0]["token_ids"]
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    text = tokenizer.decode(output_ids, skip_special_tokens=True)
+
+    print(f"Output token count: {len(output_ids)}")
+    print(f"Generated text: {text!r}")
+
+    if "red" in text.lower():
+        print("\nModel correctly identified the red image.")
+    else:
+        print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
+
+
+if __name__ == "__main__":
+    main()

requirements/rocm.txt

@@ -22,4 +22,4 @@ timm>=1.0.17
 # To be consistent with test_quark.py
 amd-quark>=0.8.99
 # Required for faster safetensors model loading
-fastsafetensors >= 0.2.2
+fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2

requirements/test/rocm.in

@@ -55,7 +55,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
 numba==0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs,azure]==0.15.7
-fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
+fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@0.2.2 # PyPI only ships CUDA wheels
 instanttensor>=0.1.5
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0

requirements/test/rocm.txt

@@ -277,7 +277,7 @@ fastar==0.10.0
     # via fastapi-cloud-cli
 fastparquet==2026.3.0
     # via genai-perf
-fastsafetensors==0.2.2
+fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@65d80088fca7a8f567fba30415fbcc80f7d2259c
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in

tests/entrypoints/openai/test_mm_serde.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Roundtrip tests for multimodal serde used by the disagg generate endpoint."""
+
+import torch
+
+from vllm.entrypoints.serve.disagg.mm_serde import (
+    decode_mm_kwargs_item,
+    encode_mm_kwargs_item,
+)
+from vllm.entrypoints.serve.disagg.protocol import (
+    MultiModalFeatures,
+    PlaceholderRangeInfo,
+)
+from vllm.multimodal.inputs import (
+    MultiModalBatchedField,
+    MultiModalFieldElem,
+    MultiModalFlatField,
+    MultiModalKwargsItem,
+    MultiModalSharedField,
+)
+
+
+def test_mm_kwargs_item_roundtrip():
+    """Full roundtrip test with all three field types and multiple dtypes."""
+    e1 = MultiModalFieldElem(
+        data=torch.zeros(1000, dtype=torch.bfloat16),
+        field=MultiModalBatchedField(),
+    )
+    e2 = MultiModalFieldElem(
+        data=torch.ones(100, dtype=torch.int32),
+        field=MultiModalSharedField(batch_size=4),
+    )
+    e3 = MultiModalFieldElem(
+        data=torch.randn(20, dtype=torch.float32),
+        field=MultiModalFlatField(slices=[slice(0, 10), slice(10, 20)], dim=0),
+    )
+
+    item = MultiModalKwargsItem({"pixel_values": e1, "grid_thw": e2, "embeds": e3})
+    encoded = encode_mm_kwargs_item(item)
+
+    # Encoded result is a base64 string
+    assert isinstance(encoded, str)
+
+    decoded = decode_mm_kwargs_item(encoded)
+
+    assert set(decoded.keys()) == {"pixel_values", "grid_thw", "embeds"}
+    assert torch.equal(item["pixel_values"].data, decoded["pixel_values"].data)
+    assert torch.equal(item["grid_thw"].data, decoded["grid_thw"].data)
+    assert torch.equal(item["embeds"].data, decoded["embeds"].data)
+    assert isinstance(decoded["pixel_values"].field, MultiModalBatchedField)
+    assert isinstance(decoded["grid_thw"].field, MultiModalSharedField)
+    assert isinstance(decoded["embeds"].field, MultiModalFlatField)
+
+
+def test_mm_kwargs_item_none_data():
+    """Roundtrip with None data field."""
+    elem = MultiModalFieldElem(
+        data=None,
+        field=MultiModalSharedField(batch_size=2),
+    )
+    item = MultiModalKwargsItem({"empty": elem})
+    encoded = encode_mm_kwargs_item(item)
+    decoded = decode_mm_kwargs_item(encoded)
+
+    assert decoded["empty"].data is None
+    assert isinstance(decoded["empty"].field, MultiModalSharedField)
+
+
+def test_mm_kwargs_item_nested_tensors():
+    """Roundtrip with nested tensor data."""
+    nested = [torch.randn(3, 4), torch.randn(5, 4)]
+    elem = MultiModalFieldElem(
+        data=nested,
+        field=MultiModalBatchedField(),
+    )
+    item = MultiModalKwargsItem({"nested": elem})
+    encoded = encode_mm_kwargs_item(item)
+    decoded = decode_mm_kwargs_item(encoded)
+
+    decoded_data = decoded["nested"].data
+    assert len(decoded_data) == 2
+    assert torch.equal(nested[0], decoded_data[0])
+    assert torch.equal(nested[1], decoded_data[1])
+
+
+def test_mm_features_with_kwargs_data():
+    """Test that MultiModalFeatures can carry serialized tensor data."""
+    elem = MultiModalFieldElem(
+        data=torch.randn(5, 3, dtype=torch.float32),
+        field=MultiModalBatchedField(),
+    )
+    item = MultiModalKwargsItem({"pixel_values": elem})
+    encoded = encode_mm_kwargs_item(item)
+
+    features = MultiModalFeatures(
+        mm_hashes={"image": ["abc123"]},
+        mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=10)]},
+        kwargs_data={"image": [encoded]},
+    )
+
+    # JSON roundtrip
+    json_str = features.model_dump_json()
+    features2 = MultiModalFeatures.model_validate_json(json_str)
+
+    assert features2.mm_hashes == {"image": ["abc123"]}
+    assert features2.kwargs_data is not None
+    assert len(features2.kwargs_data["image"]) == 1
+
+    decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0])
+    assert torch.equal(elem.data, decoded["pixel_values"].data)

tests/entrypoints/serve/disagg/test_serving_multimodal_tokens.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for multimodal features through the /inference/v1/generate endpoint.
+
+Mirrors test_serving_tokens.py but exercises the multimodal piping
+using Qwen/Qwen3-VL-2B-Instruct end-to-end via the server's /render ->
+/generate -> /detokenize path. Intentionally avoids running the HF
+processor in the pytest parent process to keep os.fork() in sibling
+tests (e.g. test_weight_transfer_llm.py) deadlock-free.
+"""
+
+import os
+
+import httpx
+import pytest
+import pytest_asyncio
+from PIL import Image
+
+from tests.utils import RemoteOpenAIServer
+from vllm.multimodal.utils import encode_image_url
+
+MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
+GEN_ENDPOINT = "/inference/v1/generate"
+RENDER_ENDPOINT = "/v1/chat/completions/render"
+DETOKENIZE_ENDPOINT = "/detokenize"
+
+
+@pytest.fixture(scope="module")
+def test_image():
+    return Image.new("RGB", (224, 224), color=(255, 0, 0))
+
+
+@pytest.fixture(scope="module")
+def server():
+    args = [
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        "4096",
+        "--enforce-eager",
+        "--no-enable-prefix-caching",
+    ]
+
+    envs = os.environ.copy()
+    envs["VLLM_ROCM_USE_SKINNY_GEMM"] = "0"
+
+    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=envs) as remote_server:
+        yield remote_server
+
+
+@pytest_asyncio.fixture
+async def client(server: RemoteOpenAIServer):
+    transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
+    headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
+    async with httpx.AsyncClient(
+        transport=transport,
+        base_url=server.url_root,
+        timeout=600,
+        headers=headers,
+    ) as c:
+        yield c
+
+
+@pytest.mark.asyncio
+async def test_render_to_generate_roundtrip(client, test_image):
+    """End-to-end: render a multimodal chat -> feed into generate -> decode.
+
+    All preprocessing and detokenization happens in the server subprocess;
+    the pytest parent never imports transformers or touches torch tensors.
+    """
+    data_url = encode_image_url(test_image, format="PNG")
+
+    render_payload = {
+        "model": MODEL_NAME,
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image_url", "image_url": {"url": data_url}},
+                    {
+                        "type": "text",
+                        "text": "What color is this image? Answer in one word.",
+                    },
+                ],
+            }
+        ],
+    }
+
+    render_resp = await client.post(RENDER_ENDPOINT, json=render_payload)
+    render_resp.raise_for_status()
+    render_data = render_resp.json()
+
+    # Validate render output structure: keys exist and values are non-empty
+    # and well-typed.
+    assert "token_ids" in render_data
+    assert isinstance(render_data["token_ids"], list)
+    assert len(render_data["token_ids"]) > 0
+    assert all(isinstance(t, int) for t in render_data["token_ids"])
+
+    assert "features" in render_data
+    features = render_data["features"]
+    assert features is not None
+    assert isinstance(features, dict)
+
+    assert "mm_hashes" in features
+    assert "image" in features["mm_hashes"]
+    image_hashes = features["mm_hashes"]["image"]
+    assert isinstance(image_hashes, list)
+    assert len(image_hashes) > 0
+    assert all(isinstance(h, str) and h for h in image_hashes)
+
+    assert "mm_placeholders" in features
+    assert "image" in features["mm_placeholders"]
+    image_placeholders = features["mm_placeholders"]["image"]
+    assert isinstance(image_placeholders, list)
+    assert len(image_placeholders) > 0
+    for p in image_placeholders:
+        assert isinstance(p.get("offset"), int)
+        assert isinstance(p.get("length"), int)
+        assert p["length"] > 0
+
+    assert "kwargs_data" in features
+    assert "image" in features["kwargs_data"]
+    assert len(features["kwargs_data"]["image"]) > 0
+
+    # Build generate request from render output
+    generate_payload = render_data
+    generate_payload["sampling_params"] = {
+        "max_tokens": 10,
+        "temperature": 0.0,
+    }
+
+    gen_resp = await client.post(GEN_ENDPOINT, json=generate_payload)
+    gen_resp.raise_for_status()
+    gen_data = gen_resp.json()
+
+    assert "choices" in gen_data
+    assert isinstance(gen_data["choices"], list)
+    assert len(gen_data["choices"]) >= 1
+    choice = gen_data["choices"][0]
+    assert "token_ids" in choice
+    assert isinstance(choice["token_ids"], list)
+    assert len(choice["token_ids"]) > 0
+    assert all(isinstance(t, int) for t in choice["token_ids"])
+
+    detok_resp = await client.post(
+        DETOKENIZE_ENDPOINT,
+        json={"model": MODEL_NAME, "tokens": choice["token_ids"]},
+    )
+    detok_resp.raise_for_status()
+    detok_data = detok_resp.json()
+    assert "prompt" in detok_data
+    text = detok_data["prompt"]
+    assert isinstance(text, str)
+    assert len(text) > 0
+    assert "red" in text.lower(), (
+        f"Expected model to identify the red image, got: {text!r}"
+    )

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/moe/test_moe.py

@@ -14,8 +14,6 @@ import pytest
 import torch
 from torch.nn import Parameter
 from torch.nn import functional as F
-from transformers import MixtralConfig
-from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
 
 import vllm.model_executor.layers.fused_moe  # noqa
 from tests.kernels.moe.utils import (
@@ -24,10 +22,7 @@ from tests.kernels.moe.utils import (
     modular_triton_fused_moe,
 )
 from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe
-from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.distributed.parallel_state import init_distributed_environment
-from vllm.forward_context import get_forward_context, set_forward_context
 from vllm.model_executor.layers.fused_moe import (
     MoEActivation,
     fused_topk,
@@ -56,12 +51,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
     marlin_quantize,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
-from vllm.model_executor.models.mixtral import MixtralMoE
 from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType, scalar_types
 from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
-from vllm.v1.worker.workspace import init_workspace_manager
 
 
 def iterative_moe(
@@ -150,12 +143,14 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
     {
         "a_type": [scalar_types.bfloat16],
         "b_type": scalar_types.float4_e2m1f,
+        "c_type": [scalar_types.bfloat16],
         "group_blocks": [2],
     },
     # MXFP8
     {
         "a_type": [scalar_types.bfloat16],
         "b_type": scalar_types.float8_e4m3fn,
+        "c_type": [scalar_types.bfloat16],
         "group_blocks": [2],
     },
     # AWQ-INT4 with INT8 activation
@@ -681,154 +676,35 @@ def test_fused_moe_wn16(
     torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("padding", [True, False])
-@pytest.mark.parametrize(
-    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
-)
-@torch.inference_mode()
-def test_mixtral_moe(
-    default_vllm_config,
-    dist_init,
-    dtype: torch.dtype,
-    padding: bool,
-    use_rocm_aiter: bool,
-    monkeypatch,
-):
-    """Make sure our Mixtral MoE implementation agrees with the one from
-    huggingface."""
-
-    # Explicitly set AITER env var based on test parameter to ensure
-    # consistent behavior regardless of external environment
-    monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0")
-    rocm_aiter_ops.refresh_env_variables()
-
-    if use_rocm_aiter and dtype == torch.float32:
-        pytest.skip("AITER ROCm test skip for float32")
-
-    monkeypatch.setenv("RANK", "0")
-    monkeypatch.setenv("LOCAL_RANK", "0")
-    monkeypatch.setenv("WORLD_SIZE", "1")
-    monkeypatch.setenv("MASTER_ADDR", "localhost")
-    monkeypatch.setenv("MASTER_PORT", "12345")
-    init_distributed_environment()
-    init_workspace_manager(torch.accelerator.current_device_index())
-
-    # Instantiate our and huggingface's MoE blocks
-    vllm_config.compilation_config.static_forward_context = dict()
-    with set_current_vllm_config(vllm_config), set_forward_context(None, vllm_config):
-        config = MixtralConfig()
-        hf_moe = MixtralSparseMoeBlock(config).to(dtype).to("cuda")
-        vllm_moe = MixtralMoE(
-            num_experts=config.num_local_experts,
-            top_k=config.num_experts_per_tok,
-            hidden_size=config.hidden_size,
-            intermediate_size=config.intermediate_size,
-            params_dtype=dtype,
-            tp_size=1,
-            dp_size=1,
-        ).cuda()
-
-        # Load the weights
-        vllm_moe.gate.weight.data[:] = hf_moe.gate.weight.data
-        if isinstance(hf_moe.experts, torch.nn.ModuleList):
-            # Transformers v4
-            for i in range(config.num_local_experts):
-                weights = (
-                    hf_moe.experts[i].w1.weight.data,
-                    hf_moe.experts[i].w3.weight.data,
-                )
-                vllm_moe.experts.w13_weight[i][:] = torch.cat(weights, dim=0)
-                vllm_moe.experts.w2_weight[i][:] = hf_moe.experts[i].w2.weight.data
-        else:
-            # Transformers v5
-            vllm_moe.experts.w13_weight.data[:] = hf_moe.experts.gate_up_proj.data
-            vllm_moe.experts.w2_weight.data[:] = hf_moe.experts.down_proj.data
-            # TODO: remove this line after https://github.com/huggingface/transformers/pull/43622
-            hf_moe.experts.config._experts_implementation = "eager"
-
-        # Generate input batch of dimensions [batch_size, seq_len, hidden_dim]
-        hf_inputs = torch.randn((1, 64, config.hidden_size)).to(dtype).to("cuda")
-        # vLLM uses 1D query [num_tokens, hidden_dim]
-        vllm_inputs = hf_inputs.flatten(0, 1)
-
-        # Pad the weight if moe padding is enabled
-        if padding:
-            vllm_moe.experts.w13_weight = Parameter(
-                F.pad(vllm_moe.experts.w13_weight, (0, 128), "constant", 0)[
-                    ..., 0:-128
-                ],
-                requires_grad=False,
-            )
-            vllm_moe.experts.w2_weight = Parameter(
-                F.pad(vllm_moe.experts.w2_weight, (0, 128), "constant", 0)[..., 0:-128],
-                requires_grad=False,
-            )
-            torch.accelerator.synchronize()
-            torch.accelerator.empty_cache()
-
-        # FIXME (zyongye) fix this after we move self.kernel
-        # assignment in FusedMoE.__init__
-
-        vllm_moe.experts.quant_method.process_weights_after_loading(vllm_moe.experts)
-
-        # need to override the forward context for unittests, otherwise it assumes
-        # we're running the model forward pass (the model specified in vllm_config)
-        get_forward_context().all_moe_layers = None
-
-        # Run forward passes for both MoE blocks
-        hf_states = hf_moe.forward(hf_inputs)
-        if isinstance(hf_states, tuple):
-            # Transformers v4
-            hf_states = hf_states[0]
-        vllm_states = vllm_moe.forward(vllm_inputs)
-
-    mixtral_moe_tol = {
-        torch.float32: 1e-3,
-        torch.float16: 1e-3,
-        torch.bfloat16: 1e-2,
-    }
-
-    if use_rocm_aiter:
-        # The values of rtol and atol are set based on the tests in ROCM AITER package.
-        # https://github.com/ROCm/aiter/blob/dfed377f4be7da96ca2d75ac0761f569676f7240/op_tests/test_moe.py#L174
-        torch.testing.assert_close(
-            hf_states.flatten(0, 1), vllm_states, rtol=0.01, atol=100
-        )
-    else:
-        torch.testing.assert_close(
-            hf_states.flatten(0, 1),
-            vllm_states,
-            rtol=mixtral_moe_tol[dtype],
-            atol=mixtral_moe_tol[dtype],
-        )
+MARLIN_MOE_SCENARIOS = [
+    # (m, n, k, e, topk, ep_size, act_order, is_k_full)
+    # No act_order: is_k_full=True matches usual case (marlin_is_k_full).
+    # N>=256 required for Marlin kernel thread config for MXFP8.
+    # Single token, small matrices
+    (1, 128, 256, 5, 2, 1, False, True),
+    # Single token, large matrices
+    (1, 1024, 2048, 5, 2, 1, False, True),
+    # Unaligned m, small matrices
+    (133, 256, 256, 5, 2, 1, False, True),
+    # Unaligned m, large matrices
+    (133, 1024, 2048, 12, 3, 1, False, True),
+    # Aligned batch, small matrices
+    (128, 256, 256, 5, 2, 1, False, True),
+    # Aligned batch, large matrices
+    (128, 1024, 2048, 12, 3, 1, False, True),
+    # Expert parallelism
+    (64, 1024, 2048, 12, 3, 4, False, True),
+    # Act order with is_k_full=True (no tensor parallelism)
+    (1, 1024, 2048, 5, 2, 1, True, True),
+    # Act order with is_k_full=False (tensor parallelism)
+    (133, 256, 256, 5, 2, 1, True, False),
+]
 
 
 def marlin_moe_generate_valid_test_cases():
     import itertools
 
-    m_list = [1, 123, 666]
-    n_list = [128, 1024]
-    k_list = [256, 2048]
-    e_list = [5, 12]
-    topk_list = [2, 3]
-    ep_size_list = [1, 4]
-    act_order_list = [True, False]
-    is_k_full_list = [True, False]
-
-    all_combinations = itertools.product(
-        MOE_MARLIN_QUANT_TEST_CONFIGS,
-        m_list,
-        n_list,
-        k_list,
-        e_list,
-        topk_list,
-        ep_size_list,
-        act_order_list,
-        is_k_full_list,
-    )
-
-    def is_invalid(
+    def is_valid(
         a_type,
         b_type,
         c_type,
@@ -845,39 +721,42 @@ def marlin_moe_generate_valid_test_cases():
         group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
         if group_size > 0 and k % group_size != 0:
             return False
-
         if act_order and group_size in [-1, k, n]:
             return False
         if group_size in [k, n]:
             return False
-        if not act_order and is_k_full:
+        if b_type == scalar_types.float8_e4m3fn and group_size == 32 and is_k_full:
             return False
-
         return a_type.size_bits < 16 or a_type is c_type
 
     cases = []
-    for case in all_combinations:
-        quant_test_config, m, n, k, _, _, _, act_order, *_ = case
-        if act_order and not quant_test_config.get("support_act_order", False):
-            continue
-
+    for quant_test_config in MOE_MARLIN_QUANT_TEST_CONFIGS:
         f16_types = [scalar_types.float16]
-        inner_combinations = itertools.product(
-            quant_test_config.get("a_type", f16_types),
-            [quant_test_config["b_type"]],
-            quant_test_config.get("c_type", f16_types),
-            quant_test_config["group_blocks"],
+        inner_combinations = list(
+            itertools.product(
+                quant_test_config.get("a_type", f16_types),
+                [quant_test_config["b_type"]],
+                quant_test_config.get("c_type", f16_types),
+                quant_test_config["group_blocks"],
+            )
         )
 
+        supports_act_order = quant_test_config.get("support_act_order", False)
+
         for sub_case in inner_combinations:
             if (
                 sub_case[0] == scalar_types.float8_e4m3fn
                 and current_platform.get_device_capability() not in [89, 120]
             ):
                 continue
-            args = sub_case + (m, n, k) + case[4:]
-            if is_invalid(*args):
-                cases.append(args)
+
+            for scenario in MARLIN_MOE_SCENARIOS:
+                m, n, k, e, topk, ep_size, act_order, is_k_full = scenario
+                if act_order and not supports_act_order:
+                    continue
+                args = sub_case + (m, n, k, e, topk, ep_size, act_order, is_k_full)
+                if is_valid(*args):
+                    cases.append(args)
     return cases
 
 

tests/kernels/moe/test_mxfp4_moe.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Tests for SM100 CUTLASS MXFP4 x MXFP4 grouped MoE kernels."""
+
+import random
+
+import pytest
+import torch
+
+from tests.kernels.utils import torch_moe_single
+from vllm import _custom_ops as ops
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import set_random_seed
+
+random.seed(42)
+set_random_seed(42)
+
+MXFP4_BLOCK_SIZE = 32
+
+
+def align(val: int, alignment: int = 128) -> int:
+    return int((val + alignment - 1) // alignment * alignment)
+
+
+def calc_diff(x, y):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def is_sm100_supported() -> bool:
+    return current_platform.is_cuda() and current_platform.is_device_capability_family(
+        100
+    )
+
+
+def compute_ref_output(
+    input_tensor: torch.Tensor,
+    weight_list: list[torch.Tensor],
+    expert_offsets: list[int],
+    expert_offset: int,
+    num_experts: int,
+) -> torch.Tensor:
+    """Reference output using torch_moe_single with top-1 routing."""
+    score = torch.full(
+        (expert_offset, num_experts),
+        -1e9,
+        device=input_tensor.device,
+        dtype=torch.float32,
+    )
+    for g in range(num_experts):
+        start = expert_offsets[g]
+        end = expert_offsets[g + 1] if g + 1 < num_experts else expert_offset
+        score[start:end, g] = 0.0
+
+    return torch_moe_single(
+        input_tensor, torch.stack(weight_list, dim=0), score, topk=1
+    )
+
+
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason="cutlass_mxfp4_group_mm requires CUDA SM100",
+)
+@pytest.mark.parametrize("num_experts", [8, 16, 32])
+@pytest.mark.parametrize("out_dtype", [torch.bfloat16])
+def test_cutlass_mxfp4_grouped_mm(num_experts, out_dtype):
+    """
+    Test the MXFP4 grouped GEMM kernel by:
+    1. Creating random per-expert inputs and weights
+    2. Quantizing both to MXFP4 using the CUDA kernel
+    3. Running the CUTLASS grouped GEMM
+    4. Comparing against BF16 reference
+    """
+    device = "cuda"
+    alignment = 128
+    # N and K must be multiples of 128 for clean swizzle layout
+    n_g = random.randint(1, 16) * alignment
+    k_g = random.randint(1, 16) * alignment
+
+    expert_offset = 0
+    expert_offsets_input = []
+    problem_sizes = []
+    input_list = []
+    weight_list = []
+
+    for g in range(num_experts):
+        m_g = random.randint(1, 256)
+        expert_offsets_input.append(expert_offset)
+        expert_offset += m_g
+        problem_sizes.append([m_g, n_g, k_g])
+
+        input_list.append(
+            torch.normal(0.0, std=0.5, size=(m_g, k_g), device=device, dtype=out_dtype)
+        )
+        weight_list.append(
+            torch.normal(0.0, std=0.5, size=(n_g, k_g), device=device, dtype=out_dtype)
+        )
+
+    input_tensor = torch.concat(input_list, dim=0)  # [M_total, K]
+
+    # --- Quantize INPUTS via mxfp4_experts_quant ---
+    input_bs_offsets = []
+    tot = 0
+    for g in range(num_experts):
+        input_bs_offsets.append(tot)
+        tot += align(problem_sizes[g][0], 128)
+    input_bs_offsets.append(tot)
+
+    _inp_expert_offsets = torch.tensor(
+        expert_offsets_input + [expert_offset], device=device, dtype=torch.int32
+    )
+    _inp_bs_offsets = torch.tensor(input_bs_offsets, device=device, dtype=torch.int32)
+
+    input_quant, input_sf = ops.mxfp4_experts_quant(
+        input_tensor,
+        _inp_expert_offsets,
+        _inp_bs_offsets,
+        num_experts,
+        topk=1,
+    )
+
+    # --- Quantize WEIGHTS via mxfp4_experts_quant ---
+    # Treat each expert's N weight rows as an "expert" with N tokens
+    weight_tensor = torch.concat(weight_list, dim=0)  # [E*N, K]
+    weight_expert_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
+    # N is always multiple of 128, so blockscale offsets are clean
+    weight_bs_offsets = [g * n_g for g in range(num_experts)] + [num_experts * n_g]
+
+    _wt_expert_offsets = torch.tensor(
+        weight_expert_offsets, device=device, dtype=torch.int32
+    )
+    _wt_bs_offsets = torch.tensor(weight_bs_offsets, device=device, dtype=torch.int32)
+
+    weight_quant, weight_sf = ops.mxfp4_experts_quant(
+        weight_tensor,
+        _wt_expert_offsets,
+        _wt_bs_offsets,
+        num_experts,
+        topk=1,
+    )
+
+    # Reshape weight quantized data to [E, N, K//2]
+    weight_quant = weight_quant[: num_experts * n_g].view(num_experts, n_g, k_g // 2)
+
+    # Reshape weight scale factors to [E, N, K//32]
+    # The quant kernel produces uint8 SF buffer. Each row has K//32 SFs.
+    scales_per_row = k_g // MXFP4_BLOCK_SIZE
+    weight_sf_flat = weight_sf.view(-1)[: num_experts * n_g * scales_per_row]
+    weight_sf_3d = weight_sf_flat.view(num_experts, n_g, scales_per_row)
+
+    # Output
+    output = torch.empty((expert_offset, n_g), device=device, dtype=out_dtype)
+
+    _problem_sizes = torch.tensor(problem_sizes, device=device, dtype=torch.int32)
+    _expert_offsets = torch.tensor(
+        expert_offsets_input, device=device, dtype=torch.int32
+    )
+    _input_bs = torch.tensor(input_bs_offsets[:-1], device=device, dtype=torch.int32)
+
+    # Run the MXFP4 grouped GEMM
+    ops.cutlass_mxfp4_moe_mm(
+        output,
+        input_quant,
+        weight_quant,
+        input_sf,
+        weight_sf_3d,
+        _problem_sizes,
+        _expert_offsets,
+        _input_bs,
+    )
+
+    # Reference: BF16 matmul
+    ref_output = compute_ref_output(
+        input_tensor=input_tensor,
+        weight_list=weight_list,
+        expert_offsets=expert_offsets_input,
+        expert_offset=expert_offset,
+        num_experts=num_experts,
+    )
+
+    # Compare per-expert
+    for g in range(num_experts):
+        start = expert_offsets_input[g]
+        end = expert_offsets_input[g + 1] if g + 1 < num_experts else expert_offset
+        if start == end:
+            continue
+        baseline = ref_output[start:end]
+        actual = output[start:end]
+        diff = calc_diff(actual, baseline)
+        print(
+            f"m_g={end - start} n_g={n_g} k_g={k_g} "
+            f"num_experts={num_experts}, "
+            f"out_dtype={out_dtype}, diff={diff:.5f}"
+        )
+        # FP4 quantization is very lossy (~4 bits precision)
+        # Comparing quantized vs full-precision gives cosine diff of 0.05-0.15
+        assert diff < 0.15, f"Expert {g}: diff={diff:.5f} exceeds threshold"
+
+
+@pytest.mark.skipif(
+    not is_sm100_supported(),
+    reason="mxfp4_experts_quant requires CUDA SM100",
+)
+def test_mxfp4_experts_quant_basic():
+    """
+    Basic smoke test for the MXFP4 experts quantization kernel.
+    """
+    device = "cuda"
+    num_experts = 4
+    k = 256
+    tokens_per_expert = 16
+
+    total_tokens = tokens_per_expert * num_experts
+    input_tensor = torch.randn(total_tokens, k, device=device, dtype=torch.bfloat16) / 5
+
+    expert_offsets = [i * tokens_per_expert for i in range(num_experts + 1)]
+    blockscale_offsets = [
+        align(i * tokens_per_expert, 128) for i in range(num_experts + 1)
+    ]
+
+    _expert_offsets = torch.tensor(expert_offsets, device=device, dtype=torch.int32)
+    _blockscale_offsets = torch.tensor(
+        blockscale_offsets, device=device, dtype=torch.int32
+    )
+
+    output, output_sf = ops.mxfp4_experts_quant(
+        input_tensor,
+        _expert_offsets,
+        _blockscale_offsets,
+        num_experts,
+        topk=1,
+    )
+
+    assert output.shape == (total_tokens, k // 2)
+    assert output.dtype == torch.uint8
+    assert output_sf.dtype == torch.uint8
+    assert output.any(), "Quantized output is all zeros"
+    print(
+        f"MXFP4 experts quant: output shape={output.shape}, sf shape={output_sf.shape}"
+    )
+    print("PASSED")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v", "-s"])

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/multimodal/test_inputs.py

@@ -26,11 +26,8 @@ def test_placeholder_range_get_num_embeds(is_embed, expected):
     "is_embed,expected",
     [
         (None, None),
-        (
-            torch.tensor([False, True, False, True, True]),
-            torch.tensor([0, 1, 1, 2, 3]),
-        ),
-        (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
+        (torch.tensor([False, True, False, True, True]), [0, 1, 1, 2, 3]),
+        (torch.tensor([True, True, True]), [1, 2, 3]),
     ],
 )
 def test_placeholder_range_embeds_cumsum(is_embed, expected):
@@ -41,6 +38,6 @@ def test_placeholder_range_embeds_cumsum(is_embed, expected):
         assert pr.embeds_cumsum is None
         return
 
-    assert torch.equal(pr.embeds_cumsum, expected)
+    assert pr.embeds_cumsum == expected
     # cached_property should return the same object on repeated access
     assert pr.embeds_cumsum is pr.embeds_cumsum

tests/quantization/test_turboquant.py

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.turboquant.config import (
 from vllm.model_executor.layers.quantization.turboquant.quantizer import (
     generate_wht_signs,
 )
+from vllm.platforms import current_platform
 from vllm.utils.math_utils import next_power_of_2
 
 # ============================================================================
@@ -345,7 +346,8 @@ class TestLloydMax:
 # Rotation matrix tests (GPU required)
 # ============================================================================
 
-CUDA_AVAILABLE = torch.cuda.is_available()
+GPGPU_AVAILABLE = torch.cuda.is_available() or torch.xpu.is_available()
+DEVICE_TYPE = current_platform.device_type
 
 
 def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor:
@@ -360,16 +362,16 @@ def generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Te
     return Q.to(device)
 
 
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
 class TestRotationMatrix:
     """Tests for the QR-based rotation (standalone benchmarks only)."""
 
     @pytest.mark.parametrize("dim", [64, 96, 128, 256])
     def test_rotation_matrix_shape_and_orthogonal(self, dim):
-        Pi = generate_rotation_matrix(dim, seed=42, device="cuda")
+        Pi = generate_rotation_matrix(dim, seed=42, device=DEVICE_TYPE)
         assert Pi.shape == (dim, dim)
         eye = Pi @ Pi.T
-        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
+        assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
             f"Pi not orthogonal for dim={dim}"
         )
 
@@ -385,7 +387,7 @@ class TestRotationMatrix:
 
     def test_rotation_matrix_det_is_pm1(self):
         """Orthogonal matrix determinant must be +1 or -1."""
-        Pi = generate_rotation_matrix(128, seed=42, device="cuda")
+        Pi = generate_rotation_matrix(128, seed=42, device=DEVICE_TYPE)
         det = torch.linalg.det(Pi)
         assert abs(abs(det.item()) - 1.0) < 1e-4
 
@@ -403,31 +405,31 @@ def _build_hadamard(d: int, device: str = "cpu") -> torch.Tensor:
     return (H / math.sqrt(d)).to(torch.device(device))
 
 
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
 class TestWHTRotation:
     """Tests for the WHT rotation actually used in serving."""
 
     @pytest.mark.parametrize("dim", [64, 128, 256])
     def test_wht_orthonormal(self, dim):
         """signs * H must be orthonormal: (signs*H) @ (signs*H)^T = I."""
-        signs = generate_wht_signs(dim, seed=42, device="cuda")
-        H = _build_hadamard(dim, "cuda")
+        signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE)
+        H = _build_hadamard(dim, DEVICE_TYPE)
         PiT = (signs.unsqueeze(1) * H).contiguous()
         eye = PiT @ PiT.T
-        assert torch.allclose(eye, torch.eye(dim, device="cuda"), atol=1e-5), (
+        assert torch.allclose(eye, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
             f"WHT rotation not orthonormal for dim={dim}"
         )
 
     @pytest.mark.parametrize("dim", [64, 128, 256])
     def test_wht_self_inverse(self, dim):
         """PiT should be self-inverse: PiT @ PiT = I (up to sign flip)."""
-        signs = generate_wht_signs(dim, seed=42, device="cuda")
-        H = _build_hadamard(dim, "cuda")
+        signs = generate_wht_signs(dim, seed=42, device=DEVICE_TYPE)
+        H = _build_hadamard(dim, DEVICE_TYPE)
         PiT = (signs.unsqueeze(1) * H).contiguous()
         Pi = PiT.T.contiguous()
         # Pi @ PiT should be identity (rotation then inverse)
         result = Pi @ PiT
-        assert torch.allclose(result, torch.eye(dim, device="cuda"), atol=1e-5), (
+        assert torch.allclose(result, torch.eye(dim, device=DEVICE_TYPE), atol=1e-5), (
             f"WHT rotation not self-inverse for dim={dim}"
         )
 
@@ -454,7 +456,7 @@ class TestWHTRotation:
 # ============================================================================
 
 
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="CUDA not available")
+@pytest.mark.skipif(not GPGPU_AVAILABLE, reason="GPGPU not available")
 class TestStoreDecodeRoundTrip:
     """End-to-end: store KV into TQ cache, decode, compare vs fp16 ref."""
 
@@ -487,11 +489,11 @@ class TestStoreDecodeRoundTrip:
         block_size = 16
         num_blocks = 1
 
-        device = torch.device("cuda")
+        device = torch.device(DEVICE_TYPE)
 
         # Generate rotation
         signs = generate_wht_signs(D, seed=42, device=device)
-        H = _build_hadamard(D, "cuda")
+        H = _build_hadamard(D, DEVICE_TYPE)
         PiT = (signs.unsqueeze(1) * H).contiguous().float()
         Pi = PiT.T.contiguous()
 

tests/test_regression.py

@@ -17,22 +17,6 @@ from vllm import LLM, SamplingParams
 from vllm.platforms import current_platform
 
 
-@pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len")
-def test_duplicated_ignored_sequence_group():
-    """https://github.com/vllm-project/vllm/issues/1655"""
-
-    sampling_params = SamplingParams(temperature=0.01, top_p=0.1, max_tokens=256)
-    llm = LLM(
-        model="distilbert/distilgpt2",
-        max_num_batched_tokens=4096,
-        tensor_parallel_size=1,
-    )
-    prompts = ["This is a short prompt", "This is a very long prompt " * 1000]
-    outputs = llm.generate(prompts, sampling_params=sampling_params)
-
-    assert len(prompts) == len(outputs)
-
-
 @pytest.mark.parametrize(
     "model",
     [

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/utils_/test_serial_utils.py

@@ -7,17 +7,39 @@ from tests.models.utils import check_embeddings_close
 from vllm.utils.serial_utils import (
     EMBED_DTYPES,
     ENDIANNESS,
+    MM_METADATA_DTYPES,
     EmbedDType,
     Endianness,
+    MmMetadataDType,
     binary2tensor,
     tensor2binary,
 )
 
+FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys())
+INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys())
+
+
+def _build_integer_tensor(
+    embed_dtype: MmMetadataDType, shape: tuple[int, ...]
+) -> torch.Tensor:
+    torch_dtype = MM_METADATA_DTYPES[embed_dtype].torch_dtype
+
+    if torch_dtype is torch.bool:
+        return torch.randint(0, 2, shape, dtype=torch.int32).to(torch.bool)
+    if torch_dtype is torch.uint8:
+        return torch.randint(0, 256, shape, dtype=torch.uint8)
+    if torch_dtype is torch.int32:
+        return torch.randint(-(2**20), 2**20, shape, dtype=torch.int32)
+    if torch_dtype is torch.int64:
+        return torch.randint(-(2**62), 2**62, shape, dtype=torch.int64)
+
+    raise AssertionError(f"Unsupported non-floating embed dtype: {embed_dtype}")
+
 
 @pytest.mark.parametrize("endianness", ENDIANNESS)
-@pytest.mark.parametrize("embed_dtype", EMBED_DTYPES.keys())
+@pytest.mark.parametrize("embed_dtype", FLOAT_EMBED_DTYPES)
 @torch.inference_mode()
-def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
+def test_encode_and_decode_floats(embed_dtype: EmbedDType, endianness: Endianness):
     for i in range(10):
         tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
         shape = tensor.shape
@@ -40,3 +62,20 @@ def test_encode_and_decode(embed_dtype: EmbedDType, endianness: Endianness):
             name_1="new",
             tol=1e-2,
         )
+
+
+@pytest.mark.parametrize("endianness", ENDIANNESS)
+@pytest.mark.parametrize("embed_dtype", INTEGER_EMBED_DTYPES)
+@torch.inference_mode()
+def test_encode_and_decode_integers(
+    embed_dtype: MmMetadataDType, endianness: Endianness
+):
+    shape = (2, 3, 5, 7, 11, 13)
+
+    for i in range(10):
+        tensor = _build_integer_tensor(embed_dtype, shape)
+        binary = tensor2binary(tensor, embed_dtype, endianness)
+        new_tensor = binary2tensor(binary, shape, embed_dtype, endianness)
+
+        assert new_tensor.dtype == MM_METADATA_DTYPES[embed_dtype].torch_dtype
+        torch.testing.assert_close(tensor, new_tensor, atol=0, rtol=0)

vllm/_custom_ops.py

@@ -1150,6 +1150,38 @@ def cutlass_fp4_moe_mm(
     )
 
 
+def cutlass_mxfp4_moe_mm(
+    out_tensors: torch.Tensor,
+    a_tensors: torch.Tensor,
+    b_tensors: torch.Tensor,
+    a_scales: torch.Tensor,
+    b_scales: torch.Tensor,
+    problem_sizes: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    sf_offsets: torch.Tensor,
+):
+    """
+    An MXFP4 Blockscaled Group Gemm for MoE (MXFP4 x MXFP4).
+
+    Uses mx_float4_t types with E8M0 scale factors and 32-element blocks.
+    - a/b_tensors: MXFP4 packed activations/weights (uint8, 2 E2M1 per byte)
+    - a_/b_scales: E8M0 blockscales (uint8, stored in swizzled layout)
+    - Epilogue uses scalar alpha=1, beta=0 inside the CUDA op (no global scales).
+    - expert_offsets/sf_offsets: expert boundary indices
+    - problem_sizes: (num_experts, 3) with (M, N, K) per expert
+    """
+    return torch.ops._C.cutlass_mxfp4_group_mm(
+        out_tensors,
+        a_tensors,
+        b_tensors,
+        a_scales,
+        b_scales,
+        problem_sizes,
+        expert_offsets,
+        sf_offsets,
+    )
+
+
 def mxfp8_experts_quant(
     input_tensor: torch.Tensor,
     problem_sizes: torch.Tensor,
@@ -1848,6 +1880,109 @@ def silu_and_mul_scaled_fp4_experts_quant(
     return output, output_scales
 
 
+def mxfp4_experts_quant(
+    input_tensor: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+    n_experts: int,
+    topk: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Quantize input tensor to MXFP4 for packed MoE inputs.
+    Uses 32-element blocks with E8M0 (power-of-two) scale factors.
+    MXFP4 has no global scale - only block-level E8M0 scale factors.
+
+    Args:
+        input_tensor: [m_topk, k] BF16/FP16 activations
+        expert_offsets: [n_experts+1] token boundaries per expert
+        blockscale_offsets: [n_experts+1] SF row boundaries per expert
+        n_experts: number of experts
+        topk: number of top-k experts
+    Returns:
+        output: [m_topk, k//2] packed E2M1 values (uint8)
+        output_scales: E8M0 blockscales in swizzled layout (uint8 view)
+    """
+    assert not current_platform.is_rocm()
+    assert input_tensor.ndim == 2
+
+    MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
+    m_numtopk, k = input_tensor.shape
+
+    assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
+        f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
+        f"{MAX_TOKENS_PER_EXPERT})"
+        f" for cutlass_moe_mxfp4, observed m_numtopk = {m_numtopk}. Use"
+        f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
+    )
+    scales_k = k // 32
+    padded_k = (scales_k + (4 - 1)) // 4
+
+    output = torch.empty(
+        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
+    )
+    output_scales = torch.empty(
+        MAX_TOKENS_PER_EXPERT * topk,
+        padded_k,
+        dtype=torch.int32,
+        device=input_tensor.device,
+    )
+    torch.ops._C.mxfp4_experts_quant(
+        output,
+        output_scales,
+        input_tensor,
+        expert_offsets,
+        blockscale_offsets,
+        n_experts,
+    )
+    # E8M0 SFs are stored as uint8
+    output_scales = output_scales.view(torch.uint8)
+    return output, output_scales
+
+
+def silu_and_mul_mxfp4_experts_quant(
+    input_tensor: torch.Tensor,
+    expert_offsets: torch.Tensor,
+    blockscale_offsets: torch.Tensor,
+    n_experts: int,
+    topk: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Fused SiLU+Mul+MXFP4 quantization for MoE intermediate activations.
+    MXFP4 has no global scale - only block-level E8M0 scale factors.
+    """
+    assert not current_platform.is_rocm()
+    assert input_tensor.ndim == 2
+
+    MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
+    m_numtopk, k_times_2 = input_tensor.shape
+    assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
+    k = k_times_2 // 2
+
+    assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk
+    scales_k = k // 32
+    padded_k = (scales_k + (4 - 1)) // 4
+
+    output = torch.empty(
+        m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
+    )
+    output_scales = torch.empty(
+        MAX_TOKENS_PER_EXPERT * topk,
+        padded_k,
+        dtype=torch.int32,
+        device=input_tensor.device,
+    )
+    torch.ops._C.silu_and_mul_mxfp4_experts_quant(
+        output,
+        output_scales,
+        input_tensor,
+        expert_offsets,
+        blockscale_offsets,
+        n_experts,
+    )
+    output_scales = output_scales.view(torch.uint8)
+    return output, output_scales
+
+
 # fp8
 def scaled_fp8_quant(
     input: torch.Tensor,

vllm/_xpu_ops.py

@@ -22,6 +22,23 @@ else:
     except ImportError:
         from torch.library import impl_abstract as register_fake
 
+if hasattr(torch.ops._xpu_C, "fp8_gemm"):
+
+    @register_fake("_xpu_C::fp8_gemm")
+    def _fp8_gemm_fake(
+        q_input: torch.Tensor,
+        q_weight: torch.Tensor,
+        out_dtype: torch.dtype,
+        input_scales: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        input_2d = q_input.view(-1, q_input.shape[-1])
+        M = input_2d.size(0)
+        N = q_weight.size(1)
+        return torch.empty((M, N), dtype=out_dtype, device=q_input.device)
+
+
 if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"):
 
     @register_fake("_xpu_C::fp8_gemm_w8a16")

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/distributed/device_communicators/xpu_communicator.py

@@ -47,9 +47,10 @@ class XpuCommunicator(DeviceCommunicatorBase):
                 self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
                 logger.info("Using AgRs manager on XPU device.")
 
-    def all_reduce(self, input_) -> torch.Tensor:
-        dist.all_reduce(input_, group=self.device_group)
-        return input_
+    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
+        output = input_.clone() if torch.compiler.is_compiling() else input_
+        dist.all_reduce(output, group=self.device_group)
+        return output
 
     def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
         world_size = self.world_size

vllm/entrypoints/anthropic/serving.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 # Adapted from
-# https://github.com/vllm/vllm/entrypoints/openai/serving_chat.py
+# https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/chat_completion/serving.py
 
 """Anthropic Messages API serving handler"""
 

vllm/entrypoints/openai/chat_completion/serving.py

@@ -557,6 +557,20 @@ class OpenAIServingChat(OpenAIServing):
             and self._should_stream_with_auto_tool_parsing(request)
         )
 
+        # Determine whether required/named tool_choice should fall back to
+        # the auto tool_parser path instead of the standard JSON-based parsing.
+        # This happens when the parser declares supports_required_and_named=False
+        # (e.g. GLM models that output XML instead of JSON).
+        tool_choice_uses_parser = (
+            self.tool_parser is not None
+            and not self.tool_parser.supports_required_and_named
+            and request.tools
+            and (
+                request.tool_choice == "required"
+                or isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
+            )
+        )
+
         all_previous_token_ids: list[list[int]] | None
         function_name_returned = [False] * num_choices
         if self.tool_call_id_type == "kimi_k2":
@@ -569,7 +583,12 @@ class OpenAIServingChat(OpenAIServing):
 
         # Only one of these will be used, thus previous_texts and
         # all_previous_token_ids will not be used twice in the same iteration.
-        if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
+        if (
+            is_mistral_grammar_path
+            or tool_choice_auto
+            or tool_choice_uses_parser
+            or reasoning_parser
+        ):
             # These are only required in "auto" tool choice case
             all_previous_token_ids = [[] for _ in range(num_choices)]
             reasoning_end_arr = [False] * num_choices
@@ -764,7 +783,12 @@ class OpenAIServingChat(OpenAIServing):
                     delta_message: DeltaMessage | None
 
                     # just update previous_texts and previous_token_ids
-                    if is_mistral_grammar_path or tool_choice_auto or reasoning_parser:
+                    if (
+                        is_mistral_grammar_path
+                        or tool_choice_auto
+                        or tool_choice_uses_parser
+                        or reasoning_parser
+                    ):
                         assert previous_texts is not None
                         assert all_previous_token_ids is not None
                         previous_text = previous_texts[i]
@@ -813,7 +837,9 @@ class OpenAIServingChat(OpenAIServing):
                         if result.tools_called:
                             tools_streamed[i] = True
                     # handle streaming deltas for tools with named tool_choice
-                    elif tool_choice_function_name:
+                    # Skip when tool_choice_uses_parser so it falls through
+                    # to the auto tool_parser branches below.
+                    elif tool_choice_function_name and not tool_choice_uses_parser:
                         # When encountering think end id in prompt_token_ids
                         # i.e {"enable_thinking": False},
                         # check BEFORE calling the parser to avoid a spurious
@@ -851,7 +877,6 @@ class OpenAIServingChat(OpenAIServing):
                             ):
                                 reasoning_end_arr[i] = True
                                 if delta_message and delta_message.content:
-                                    # This need to be added to next `delta_text`
                                     current_text = delta_message.content
                                     delta_message.content = None
                                 else:
@@ -896,7 +921,12 @@ class OpenAIServingChat(OpenAIServing):
                             )
                             tools_streamed[i] = True
 
-                    elif request.tool_choice == "required":
+                    # Skip when tool_choice_uses_parser so it falls through
+                    # to the auto tool_parser branches below.
+                    elif (
+                        request.tool_choice == "required"
+                        and not tool_choice_uses_parser
+                    ):
                         assert previous_texts is not None
                         previous_text = previous_texts[i]
                         current_text = previous_text + delta_text
@@ -966,7 +996,10 @@ class OpenAIServingChat(OpenAIServing):
 
                     # update the previous values for the next iteration
                     if (
-                        is_mistral_grammar_path or tool_choice_auto or reasoning_parser
+                        is_mistral_grammar_path
+                        or tool_choice_auto
+                        or tool_choice_uses_parser
+                        or reasoning_parser
                     ) and not self.use_harmony:
                         assert previous_texts is not None
                         assert all_previous_token_ids is not None

vllm/entrypoints/openai/engine/serving.py

@@ -627,7 +627,7 @@ class OpenAIServing:
             and isinstance(request.tool_choice, ToolChoiceFunction)
         ):
             assert content is not None
-            # Forced Function Call
+            # Forced Function Call (Responses API)
             function_calls.append(
                 FunctionCall(name=request.tool_choice.name, arguments=content)
             )
@@ -636,14 +636,20 @@ class OpenAIServing:
             not use_mistral_tool_parser
             and request.tool_choice
             and isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
+            and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
         ):
+            # Named function with standard JSON-based parsing
             assert content is not None
-            # Forced Function Call
             function_calls.append(
                 FunctionCall(name=request.tool_choice.function.name, arguments=content)
             )
             content = None  # Clear content since tool is called.
-        elif not use_mistral_tool_parser and request.tool_choice == "required":
+        elif (
+            not use_mistral_tool_parser
+            and request.tool_choice == "required"
+            and (tool_parser_cls is None or tool_parser_cls.supports_required_and_named)
+        ):
+            # "required" with standard JSON-based parsing
             tool_calls = []
             with contextlib.suppress(ValidationError):
                 content = content or ""
@@ -662,15 +668,30 @@ class OpenAIServing:
             use_mistral_tool_parser
             or (
                 enable_auto_tools
-                and (request.tool_choice == "auto" or request.tool_choice is None)
+                and (
+                    request.tool_choice == "auto"
+                    or request.tool_choice is None
+                    or (
+                        not tool_parser_cls.supports_required_and_named
+                        and request.tools
+                        and (
+                            request.tool_choice == "required"
+                            or isinstance(
+                                request.tool_choice,
+                                ChatCompletionNamedToolChoiceParam,
+                            )
+                        )
+                    )
+                )
             )
         ):
+            # Automatic Tool Call Parsing (also used as fallback for
+            # required/named when supports_required_and_named=False)
             if tokenizer is None:
                 raise ValueError(
                     "Tokenizer not available when `skip_tokenizer_init=True`"
                 )
 
-            # Automatic Tool Call Parsing
             try:
                 tool_parser = tool_parser_cls(tokenizer, request.tools)
             except RuntimeError as e:

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/entrypoints/serve/disagg/mm_serde.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Encode/decode utilities for multimodal tensors and field metadata
+over JSON/HTTP, used by the disaggregated generate endpoint."""
+
+from __future__ import annotations
+
+import pybase64
+
+from vllm.multimodal.inputs import MultiModalKwargsItem
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
+
+_encoder = MsgpackEncoder(size_threshold=2**62)  # force all tensors inline
+_decoder = MsgpackDecoder(t=MultiModalKwargsItem)
+
+
+def encode_mm_kwargs_item(item: MultiModalKwargsItem) -> str:
+    """Serialize a MultiModalKwargsItem to a base64 string."""
+    bufs = _encoder.encode(item)
+    assert len(bufs) == 1, "All tensors should be inline"
+    return pybase64.b64encode(bufs[0]).decode("ascii")
+
+
+def decode_mm_kwargs_item(data: str) -> MultiModalKwargsItem:
+    """Deserialize a base64 string back to a MultiModalKwargsItem."""
+    raw = pybase64.b64decode(data)
+    return _decoder.decode(raw)

vllm/entrypoints/serve/disagg/protocol.py

@@ -35,14 +35,6 @@ class MultiModalFeatures(BaseModel):
     Carries hashes (for cache lookup / identification) and placeholder
     positions so the downstream `/generate` service knows *where* in
     the token sequence each multimodal item lives.
-
-    Note:
-        Phase 1 — metadata only.
-        Phase 2 should add `mm_kwargs` (processed tensor data) using a
-        binary transport so the ``/generate` side can skip re-processing.
-        The `/generate` endpoint must also be updated to inject these
-        features into `EngineInput` before passing to
-        `InputProcessor.process_inputs`.
     """
 
     mm_hashes: dict[str, list[str]]
@@ -51,6 +43,15 @@ class MultiModalFeatures(BaseModel):
     mm_placeholders: dict[str, list[PlaceholderRangeInfo]]
     """Per-modality placeholder ranges in the token sequence."""
 
+    kwargs_data: dict[str, list[str | None]] | None = None
+    """Per-modality serialized tensor data.
+
+    Each value is a list parallel to ``mm_hashes[modality]``.  A ``str``
+    entry is a base64-encoded ``MultiModalKwargsItem``; ``None`` means
+    the item should be resolved from cache.  The entire field is
+    ``None`` for metadata-only (cache-hit) responses.
+    """
+
 
 class GenerateRequest(BaseModel):
     request_id: str = Field(

vllm/entrypoints/serve/disagg/serving.py

@@ -25,6 +25,7 @@ from vllm.entrypoints.openai.engine.protocol import (
 )
 from vllm.entrypoints.openai.engine.serving import OpenAIServing, clamp_prompt_logprobs
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item
 from vllm.entrypoints.serve.disagg.protocol import (
     GenerateRequest,
     GenerateResponse,
@@ -34,8 +35,14 @@ from vllm.entrypoints.serve.disagg.protocol import (
 )
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.utils import should_include_usage
+from vllm.inputs import EngineInput, mm_input
 from vllm.logger import init_logger
 from vllm.logprobs import Logprob
+from vllm.multimodal.inputs import (
+    MultiModalKwargsItem,
+    MultiModalKwargsItems,
+    PlaceholderRange,
+)
 from vllm.outputs import RequestOutput
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.utils.collection_utils import as_list
@@ -103,11 +110,42 @@ class ServingTokens(OpenAIServing):
         if raw_request:
             raw_request.state.request_metadata = request_metadata
 
-        (engine_input,) = await self.openai_serving_render.preprocess_completion(
-            request,
-            prompt_input=request.token_ids,
-            prompt_embeds=None,
-        )
+        engine_input: EngineInput
+        if features := request.features:
+            # Convert PlaceholderRangeInfo → PlaceholderRange per modality.
+            mm_placeholders: dict[str, list[PlaceholderRange]] = {
+                modality: [
+                    PlaceholderRange(offset=p.offset, length=p.length) for p in ranges
+                ]
+                for modality, ranges in features.mm_placeholders.items()
+            }
+
+            # Deserialize tensor data when present; None → cache hit.
+            mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {}
+            if features.kwargs_data is not None:
+                for modality, items in features.kwargs_data.items():
+                    mm_kwargs[modality] = [
+                        decode_mm_kwargs_item(item) if item is not None else None
+                        for item in items
+                    ]
+            else:
+                for modality, hashes in features.mm_hashes.items():
+                    mm_kwargs[modality] = [None] * len(hashes)
+
+            engine_input = mm_input(
+                prompt_token_ids=request.token_ids,
+                mm_kwargs=MultiModalKwargsItems(mm_kwargs),
+                mm_hashes=features.mm_hashes,
+                mm_placeholders=mm_placeholders,
+                cache_salt=request.cache_salt,
+            )
+        else:
+            (engine_input,) = await self.openai_serving_render.preprocess_completion(
+                request,
+                prompt_input=request.token_ids,
+                prompt_embeds=None,
+                skip_mm_cache=True,
+            )
 
         # Schedule the request and get the result generator.
         result_generator: AsyncGenerator[RequestOutput, None] | None = None

vllm/entrypoints/serve/render/serving.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Sequence
 from http import HTTPStatus
-from typing import Any
+from typing import Any, cast
 
 from openai_harmony import Message as OpenAIMessage
 
@@ -25,6 +25,7 @@ from vllm.entrypoints.openai.parser.harmony_utils import (
     render_for_completion,
 )
 from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+from vllm.entrypoints.serve.disagg.mm_serde import encode_mm_kwargs_item
 from vllm.entrypoints.serve.disagg.protocol import (
     GenerateRequest,
     MultiModalFeatures,
@@ -37,6 +38,7 @@ from vllm.entrypoints.utils import (
 from vllm.inputs import (
     EngineInput,
     MultiModalHashes,
+    MultiModalInput,
     MultiModalPlaceholders,
     PromptType,
     SingletonPrompt,
@@ -251,6 +253,7 @@ class OpenAIServingRender:
                 default_template_kwargs=self.default_chat_template_kwargs,
                 tool_dicts=tool_dicts,
                 tool_parser=tool_parser,
+                skip_mm_cache=True,
                 reasoning_parser=self.reasoning_parser,
             )
         else:
@@ -342,6 +345,7 @@ class OpenAIServingRender:
             request,
             prompt_input=request.prompt,
             prompt_embeds=request.prompt_embeds,
+            skip_mm_cache=True,
         )
 
         return engine_inputs
@@ -357,9 +361,10 @@ class OpenAIServingRender:
         if engine_input.get("type") != "multimodal":
             return None
 
-        # At this point engine_input is a MultiModalInputs TypedDict.
-        mm_hashes: MultiModalHashes = engine_input["mm_hashes"]  # type: ignore[typeddict-item]
-        raw_placeholders: MultiModalPlaceholders = engine_input["mm_placeholders"]  # type: ignore[typeddict-item]
+        # At this point engine_input is a MultiModalInput TypedDict.
+        mm_engine_input = cast(MultiModalInput, engine_input)
+        mm_hashes: MultiModalHashes = mm_engine_input["mm_hashes"]
+        raw_placeholders: MultiModalPlaceholders = mm_engine_input["mm_placeholders"]
 
         mm_placeholders = {
             modality: [
@@ -368,9 +373,20 @@ class OpenAIServingRender:
             for modality, ranges in raw_placeholders.items()
         }
 
+        # Serialize tensor data per modality.
+        kwargs_data: dict[str, list[str | None]] | None = None
+        if raw_mm_kwargs := mm_engine_input.get("mm_kwargs"):
+            kwargs_data = {}
+            for modality, items in raw_mm_kwargs.items():
+                kwargs_data[modality] = [
+                    encode_mm_kwargs_item(item) if item is not None else None
+                    for item in items
+                ]
+
         return MultiModalFeatures(
             mm_hashes=mm_hashes,
             mm_placeholders=mm_placeholders,
+            kwargs_data=kwargs_data,
         )
 
     def _make_request_with_harmony(

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/kernels/helion/register.py

@@ -53,14 +53,15 @@ if not has_helion():
     )
 
 import helion
-from helion._compat import requires_torch_version
 from helion.autotuner.base_search import BaseAutotuner
 from helion.runtime.config import Config
 from helion.runtime.settings import default_autotuner_fn
 
 # TODO(gmagogsfm): Remove CustomOp fallback path (_get_or_register_custom_op,
 # vllm_helion_lib, direct_register_custom_op) once vLLM requires PyTorch >= 2.11.
-_HOP_AVAILABLE = requires_torch_version("2.11")
+# FIXME(gmagogsfm): Re-enable HOP path once performance regression is fixed.
+# _HOP_AVAILABLE = requires_torch_version("2.11")
+_HOP_AVAILABLE = False
 
 if _HOP_AVAILABLE:
     from helion._compat import supports_torch_compile_fusion

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/fused_moe/config.py

@@ -762,6 +762,25 @@ def nvfp4_moe_quant_config(
     )
 
 
+def mxfp4_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for MXFP4 x MXFP4 MoE.
+    MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
+    separate alphas / global activation scales in this config.
+    """
+    return FusedMoEQuantConfig.make(
+        "mxfp4",
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        per_act_token_quant=False,
+        per_out_ch_quant=False,
+        block_shape=None,
+    )
+
+
 def nvfp4_w4a16_moe_quant_config(
     g1_alphas: torch.Tensor,
     g2_alphas: torch.Tensor,

vllm/model_executor/layers/fused_moe/cutlass_moe.py

@@ -36,6 +36,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8DynamicTokenSym,
     kFp8StaticChannelSym,
     kFp8StaticTensorSym,
+    kMxfp4Dynamic,
+    kMxfp4Static,
     kNvfp4Dynamic,
     kNvfp4Static,
 )
@@ -795,6 +797,299 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
         )
 
 
+def run_cutlass_moe_mxfp4(
+    output: torch.Tensor,
+    a: torch.Tensor,
+    w1_fp4: torch.Tensor,
+    w1_blockscale: torch.Tensor,
+    w2_fp4: torch.Tensor,
+    w2_blockscale: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    activation: MoEActivation,
+    workspace13: torch.Tensor,
+    workspace2: torch.Tensor,
+    m: int,
+    n: int,
+    k: int,
+    e: int,
+    device: torch.device,
+    apply_router_weight_on_input: bool = False,
+) -> None:
+    """MXFP4 x MXFP4 MoE implementation using CUTLASS grouped GEMM."""
+    is_gated = activation.is_gated
+    w1_n = n * 2 if is_gated else n
+
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8"
+    assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8"
+    assert (
+        w1_fp4.ndim == 3
+        and w2_fp4.ndim == 3
+        and w1_blockscale.ndim == 3
+        and w2_blockscale.ndim == 3
+    ), "All Weights must be of rank 3 for cutlass_moe_mxfp4"
+    m_a, k_a = a.shape
+    e_w1, w1_n_actual, half_k_w1 = w1_fp4.shape
+    e_w2, k_w2, half_n_w2 = w2_fp4.shape
+
+    assert e_w1 == e_w2 and e_w1 == e
+    assert k_a == half_k_w1 * 2 and k == k_w2
+    assert w1_n_actual == w1_n and half_n_w2 * 2 == n
+    assert m == m_a
+    assert 2 * half_k_w1 == k_w2
+    assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
+    assert topk_weights.size(0) == m and topk_ids.size(0) == m
+
+    topk = topk_ids.size(1)
+    out_dtype = a.dtype
+    num_topk = topk_ids.size(1)
+
+    expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
+    blockscale_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
+    problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device)
+    problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device)
+
+    a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
+    c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
+
+    if apply_router_weight_on_input:
+        assert num_topk == 1, (
+            "apply_router_weight_on_input is only implemented for topk=1"
+        )
+        a.mul_(topk_weights.to(out_dtype))
+
+    ops.get_cutlass_moe_mm_data(
+        topk_ids,
+        expert_offsets,
+        problem_sizes1,
+        problem_sizes2,
+        a_map,
+        c_map,
+        e,
+        n,
+        k,
+        blockscale_offsets,
+        is_gated=is_gated,
+    )
+
+    a = ops.shuffle_rows(a, a_map)
+    rep_a_fp4, rep_a_blockscale = ops.mxfp4_experts_quant(
+        a,
+        expert_offsets,
+        blockscale_offsets,
+        e,
+        num_topk,
+    )
+    c1 = _resize_cache(workspace13, (m * topk, w1_n))
+    c2 = _resize_cache(workspace2, (m * topk, n))
+    c3 = _resize_cache(workspace13, (m * topk, k))
+
+    ops.cutlass_mxfp4_moe_mm(
+        c1,
+        rep_a_fp4,
+        w1_fp4,
+        rep_a_blockscale,
+        w1_blockscale,
+        problem_sizes1,
+        expert_offsets[:-1],
+        blockscale_offsets[:-1],
+    )
+    del rep_a_fp4, rep_a_blockscale
+    if activation == MoEActivation.SILU:
+        int_fp4, int_blockscale = ops.silu_and_mul_mxfp4_experts_quant(
+            c1, expert_offsets, blockscale_offsets, e, num_topk
+        )
+    else:
+        apply_moe_activation(activation, c2, c1)
+        int_fp4, int_blockscale = ops.mxfp4_experts_quant(
+            c2, expert_offsets, blockscale_offsets, e, num_topk
+        )
+
+    ops.cutlass_mxfp4_moe_mm(
+        c3,
+        int_fp4,
+        w2_fp4,
+        int_blockscale,
+        w2_blockscale,
+        problem_sizes2,
+        expert_offsets[:-1],
+        blockscale_offsets[:-1],
+    )
+    del int_fp4, int_blockscale
+
+    c3 = ops.shuffle_rows(c3, c_map)
+
+    assert output.dtype == out_dtype
+    if not apply_router_weight_on_input:
+        output.copy_(
+            (
+                c3.view(m, num_topk, k)
+                * topk_weights.view(m, num_topk, 1).to(out_dtype)
+            ).sum(dim=1),
+            non_blocking=True,
+        )
+    else:
+        output.copy_(c3.view(m, num_topk, k).sum(dim=1), non_blocking=True)
+    return
+
+
+def swizzle_mxfp4_scales(
+    scales: torch.Tensor,
+    N: int,
+    K: int,
+) -> torch.Tensor:
+    """Swizzle flat [N, K//32] E8M0 scales to CUTLASS tiled layout.
+
+    CUTLASS expects MX scale factors in a tiled layout:
+        [numMTiles, numKTiles, 32, 4, 4]
+    where numMTiles = ceil(N/128), numKTiles = ceil(K/128),
+    and the inner dimensions correspond to the swizzle pattern:
+        mTileIdx = mIdx / 128
+        outerMIdx = mIdx % 32
+        innerMIdx = (mIdx / 32) % 4
+        kTileIdx = kIdx / 4
+        innerKIdx = kIdx % 4
+    with kIdx = col_in_scale_space (i.e., index into K//32).
+    """
+    assert scales.dtype == torch.uint8
+    num_scale_cols = K // 32  # number of E8M0 scale values per row
+
+    num_m_tiles = (N + 127) // 128
+    num_k_tiles = (num_scale_cols + 3) // 4
+
+    # Pad N to multiple of 128 and scale_cols to multiple of 4
+    padded_N = num_m_tiles * 128
+    padded_scale_cols = num_k_tiles * 4
+
+    # Start with flat scales, pad if needed
+    padded = torch.zeros(
+        padded_N, padded_scale_cols, dtype=torch.uint8, device=scales.device
+    )
+    padded[:N, :num_scale_cols] = scales
+
+    # Reshape to tile structure:
+    # [numMTiles, 4, 32, numKTiles, 4]
+    #  mTileIdx, innerMIdx, outerMIdx, kTileIdx, innerKIdx
+    tiled = padded.reshape(num_m_tiles, 4, 32, num_k_tiles, 4)
+    # Permute to [numMTiles, numKTiles, 32, 4, 4]
+    #            (outerMIdx, innerMIdx, innerKIdx)
+    tiled = tiled.permute(0, 3, 2, 1, 4).contiguous()
+    return tiled.reshape(-1)
+
+
+class CutlassExpertsMxfp4(mk.FusedMoEExpertsModular):
+    """CUTLASS MXFP4 x MXFP4 fused MoE expert implementation."""
+
+    @property
+    def expects_unquantized_inputs(self) -> bool:
+        return True
+
+    @staticmethod
+    def _supports_current_device() -> bool:
+        p = current_platform
+        return p.is_cuda() and p.is_device_capability_family(100)
+
+    @staticmethod
+    def _supports_no_act_and_mul() -> bool:
+        return True
+
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        return (weight_key, activation_key) == (kMxfp4Static, kMxfp4Dynamic)
+
+    @staticmethod
+    def _supports_activation(activation: MoEActivation) -> bool:
+        return activation in [
+            MoEActivation.SILU,
+            MoEActivation.GELU,
+            MoEActivation.SWIGLUOAI,
+            MoEActivation.SWIGLUSTEP,
+            MoEActivation.SILU_NO_MUL,
+            MoEActivation.GELU_NO_MUL,
+            MoEActivation.RELU2_NO_MUL,
+        ]
+
+    @staticmethod
+    def _supports_parallel_config(
+        moe_parallel_config: FusedMoEParallelConfig,
+    ) -> bool:
+        return moe_parallel_config.ep_size == 1
+
+    @staticmethod
+    def activation_format() -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def supports_expert_map(self) -> bool:
+        return False
+
+    def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
+        return TopKWeightAndReduceNoOP()
+
+    def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
+        return act_dtype
+
+    def workspace_shapes(
+        self,
+        M: int,
+        N: int,
+        K: int,
+        topk: int,
+        global_num_experts: int,
+        local_num_experts: int,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        activation: MoEActivation,
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        workspace1 = (M * topk, max(2 * N, K))
+        workspace2 = (M * topk, N)
+        output = (M, K)
+        return (workspace1, workspace2, output)
+
+    def apply(
+        self,
+        output: torch.Tensor,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: MoEActivation,
+        global_num_experts: int,
+        expert_map: torch.Tensor | None,
+        a1q_scale: torch.Tensor | None,
+        a2_scale: torch.Tensor | None,
+        workspace13: torch.Tensor | None,
+        workspace2: torch.Tensor | None,
+        expert_tokens_meta: mk.ExpertTokensMetadata | None,
+        apply_router_weight_on_input: bool,
+    ):
+        e, m, n, k, _ = self.moe_problem_size(hidden_states, w1, w2, topk_ids)
+        n = w2.shape[2] * 2
+
+        run_cutlass_moe_mxfp4(
+            output=output,
+            a=hidden_states,
+            w1_fp4=w1,
+            w1_blockscale=self.w1_scale,
+            w2_fp4=w2,
+            w2_blockscale=self.w2_scale,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            workspace13=workspace13,
+            workspace2=workspace2,
+            m=m,
+            n=n,
+            k=k,
+            e=e,
+            device=hidden_states.device,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+        )
+
+
 # W4A8
 def run_cutlass_moe_w4a8_fp8(
     output: torch.Tensor,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py

@@ -4,6 +4,7 @@
 
 import torch
 
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
@@ -11,6 +12,10 @@ from vllm.model_executor.layers.fused_moe import (
 )
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEQuantConfig,
+    mxfp4_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.cutlass_moe import (
+    CutlassExpertsMxfp4,
 )
 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
     MarlinExperts,
@@ -36,7 +41,14 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
         super().__init__(moe)
         self.group_size = 32
         self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
-        self.experts_cls = MarlinExperts
+        self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
+        self.experts_cls: type[mk.FusedMoEExperts]
+        if self.use_cutlass_mxfp4:
+            logger.info_once("Using CutlassExpertsMxfp4 for MXFP4 MoE", scope="local")
+            self.experts_cls = CutlassExpertsMxfp4
+        else:
+            logger.info_once("Using MarlinExperts for MXFP4 MoE", scope="local")
+            self.experts_cls = MarlinExperts
 
     def create_weights(
         self,
@@ -109,11 +121,19 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
     def get_fused_moe_quant_config(
         self, layer: torch.nn.Module
     ) -> FusedMoEQuantConfig | None:
-        return make_mxfp4_moe_quant_config(
-            mxfp4_backend=self.mxfp4_backend,
-            w1_scale=layer.w13_weight_scale,
-            w2_scale=layer.w2_weight_scale,
-        )
+        if self.use_cutlass_mxfp4:
+            # W4A4: both weights and activations quantized to MXFP4
+            return mxfp4_moe_quant_config(
+                w1_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+            )
+        else:
+            # W4A16: weight-only via Marlin
+            return make_mxfp4_moe_quant_config(
+                mxfp4_backend=self.mxfp4_backend,
+                w1_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+            )
 
     def process_weights_after_loading(self, layer: FusedMoE) -> None:
         layer.w13_weight = torch.nn.Parameter(
@@ -126,13 +146,45 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
         )
         delattr(layer, "w2_weight_packed")
 
-        logger.warning_once(
-            "Your GPU does not have native support for FP4 computation but "
-            "FP4 quantization is being used. Weight-only FP4 compression "
-            "will be used leveraging the Marlin kernel. This may degrade "
-            "performance for compute-heavy workloads."
-        )
-        prepare_moe_fp4_layer_for_marlin(layer)
+        if self.use_cutlass_mxfp4:
+            # Swizzle weight scales from flat checkpoint layout [E, N, K//32]
+            # to CUTLASS tiled layout [E, numMTiles*numKTiles*512].
+            from vllm.model_executor.layers.fused_moe.cutlass_moe import (
+                swizzle_mxfp4_scales,
+            )
+
+            E = layer.w13_weight_scale.shape[0]
+            w13_N = layer.w13_weight_scale.shape[1]
+            w13_scale_K = layer.w13_weight_scale.shape[2]
+            w13_K = w13_scale_K * 32
+
+            w2_M = layer.w2_weight_scale.shape[1]
+            w2_scale_N = layer.w2_weight_scale.shape[2]
+            w2_N = w2_scale_N * 32
+
+            swizzled_w13 = []
+            swizzled_w2 = []
+            for e_idx in range(E):
+                s13 = layer.w13_weight_scale[e_idx]
+                sw13 = swizzle_mxfp4_scales(s13, w13_N, w13_K)
+                swizzled_w13.append(sw13.reshape(w13_N, w13_scale_K))
+                s2 = layer.w2_weight_scale[e_idx]
+                sw2 = swizzle_mxfp4_scales(s2, w2_M, w2_N)
+                swizzled_w2.append(sw2.reshape(w2_M, w2_scale_N))
+            layer.w13_weight_scale = torch.nn.Parameter(
+                torch.stack(swizzled_w13), requires_grad=False
+            )
+            layer.w2_weight_scale = torch.nn.Parameter(
+                torch.stack(swizzled_w2), requires_grad=False
+            )
+        else:
+            logger.warning_once(
+                "Your GPU does not have native support for FP4 computation "
+                "but FP4 quantization is being used. Weight-only FP4 "
+                "compression will be used leveraging the Marlin kernel. "
+                "This may degrade performance for compute-heavy workloads."
+            )
+            prepare_moe_fp4_layer_for_marlin(layer)
 
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
         if self.moe_quant_config is not None:

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

@@ -67,6 +67,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
 from .interfaces import (
     MultiModalEmbeddings,
     SupportsEagle3,
+    SupportsLoRA,
     SupportsMultiModal,
     SupportsPP,
 )
@@ -880,6 +881,7 @@ class Gemma4ForConditionalGeneration(
     nn.Module,
     SupportsMultiModal,
     SupportsPP,
+    SupportsLoRA,
     SupportsEagle3,
 ):
     packed_modules_mapping = {
@@ -1254,9 +1256,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
@@ -1357,10 +1360,16 @@ class Gemma4ForConditionalGeneration(
 
     def get_mm_mapping(self) -> MultiModelKeys:
         """Get the module prefix mapping for multimodal models."""
+        connectors = ["embed_vision"]
+        tower_models = ["vision_tower"]
+        if self.audio_tower is not None:
+            connectors.append("embed_audio")
+            tower_models.append("audio_tower")
+
         return MultiModelKeys.from_string_field(
             language_model="language_model",
-            connector=["embed_vision", "embed_audio"],
-            tower_model=["vision_tower", "audio_tower"],
+            connector=connectors,
+            tower_model=tower_models,
         )
 
     @classmethod

vllm/model_executor/models/glmasr.py

@@ -66,7 +66,7 @@ from .interfaces import (
     SupportsTranscription,
 )
 from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
-from .whisper import ISO639_1_SUPPORTED_LANGS
+from .whisper import ISO639_1_SUPPORTED_LANGS, _create_fake_bias_for_k_proj
 
 
 class GlmAsrEncoderRotaryEmbedding(nn.Module):
@@ -499,6 +499,8 @@ class GlmAsrEncoder(nn.Module):
         """Custom weight loading to handle q_proj/k_proj/v_proj -> qkv_proj mapping."""
         from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
+        weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
+
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),

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/multimodal/inputs.py

@@ -145,14 +145,15 @@ class PlaceholderRange:
     """
 
     @cached_property
-    def embeds_cumsum(self) -> torch.Tensor | None:
-        return None if self.is_embed is None else self.is_embed.cumsum(dim=0)
+    def embeds_cumsum(self) -> list[int] | None:
+        # python list so python indexing avoids torch C++ overhead/conversions/deallocs
+        return None if self.is_embed is None else self.is_embed.cumsum(dim=0).tolist()
 
     def get_num_embeds(self) -> int:
         if self.embeds_cumsum is None:
             return self.length
 
-        return int(self.embeds_cumsum[-1])
+        return self.embeds_cumsum[-1] if self.embeds_cumsum else 0
 
     def get_embeds_indices_in_range(
         self, start_idx: int, end_idx: int
@@ -170,10 +171,8 @@ class PlaceholderRange:
         if self.embeds_cumsum is None:
             return start_idx, end_idx
 
-        embeds_start_idx = (
-            int(self.embeds_cumsum[start_idx - 1]) if start_idx > 0 else 0
-        )
-        embeds_end_idx = int(self.embeds_cumsum[end_idx - 1])
+        embeds_start_idx = self.embeds_cumsum[start_idx - 1] if start_idx > 0 else 0
+        embeds_end_idx = self.embeds_cumsum[end_idx - 1] if end_idx > 0 else 0
 
         return embeds_start_idx, embeds_end_idx
 

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

@@ -131,6 +131,7 @@ def _get_backend_priorities(
                 AttentionBackendEnum.FLASH_ATTN,
                 AttentionBackendEnum.TRITON_ATTN,
                 AttentionBackendEnum.FLEX_ATTENTION,
+                AttentionBackendEnum.TURBOQUANT,
             ]
         else:
             return [
@@ -138,6 +139,7 @@ def _get_backend_priorities(
                 AttentionBackendEnum.FLASHINFER,
                 AttentionBackendEnum.TRITON_ATTN,
                 AttentionBackendEnum.FLEX_ATTENTION,
+                AttentionBackendEnum.TURBOQUANT,
             ]
 
 
@@ -255,11 +257,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/platforms/rocm.py

@@ -382,6 +382,7 @@ def _get_backend_priorities(
     if is_aiter_found_and_supported():
         backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
     backends.append(AttentionBackendEnum.TRITON_ATTN)
+    backends.append(AttentionBackendEnum.TURBOQUANT)
 
     return backends
 

vllm/reasoning/__init__.py

@@ -60,6 +60,10 @@ _REASONING_PARSERS_TO_REGISTER = {
         "kimi_k2_reasoning_parser",
         "KimiK2ReasoningParser",
     ),
+    "mimo": (
+        "qwen3_reasoning_parser",
+        "Qwen3ReasoningParser",
+    ),
     "minimax_m2": (
         "minimax_m2_reasoning_parser",
         "MiniMaxM2ReasoningParser",

vllm/tool_parsers/__init__.py

@@ -94,6 +94,10 @@ _TOOL_PARSERS_TO_REGISTER = {
         "longcat_tool_parser",
         "LongcatFlashToolParser",
     ),
+    "mimo": (
+        "qwen3xml_tool_parser",
+        "Qwen3XMLToolParser",
+    ),
     "minimax_m2": (
         "minimax_m2_tool_parser",
         "MinimaxM2ToolParser",

vllm/tool_parsers/abstract_tool_parser.py

@@ -44,6 +44,17 @@ class ToolParser:
     derived classes.
     """
 
+    # When True (default), the serving layer uses the standard JSON-based
+    # parsing for tool_choice="required" and named function tool_choice,
+    # which works for models where guided decoding produces well-formed
+    # JSON output (e.g. Hermes).
+    # Subclasses set False when the standard parsing does not work for
+    # their model's output format (e.g. GLM models that use XML).  When
+    # False, the serving layer falls back to the tool_parser's
+    # extract_tool_calls / extract_tool_calls_streaming methods for
+    # required/named tool_choice, treating them the same as "auto".
+    supports_required_and_named: bool = True
+
     def __init__(
         self,
         tokenizer: TokenizerLike,

vllm/tool_parsers/glm47_moe_tool_parser.py

@@ -23,6 +23,8 @@ logger = init_logger(__name__)
 
 
 class Glm47MoeModelToolParser(Glm4MoeModelToolParser):
+    supports_required_and_named = False
+
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         super().__init__(tokenizer, tools)
         # GLM-4.7 format: <tool_call>func_name[<arg_key>...]*</tool_call>

vllm/tool_parsers/glm4_moe_tool_parser.py

@@ -20,6 +20,7 @@ import regex as re
 
 from vllm.entrypoints.chat_utils import make_tool_call_id
 from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionNamedToolChoiceParam,
     ChatCompletionRequest,
 )
 from vllm.entrypoints.openai.engine.protocol import (
@@ -50,6 +51,8 @@ class Glm4MoeModelToolParser(ToolParser):
     call, and diffs against what was previously sent to emit only new content.
     """
 
+    supports_required_and_named = False
+
     def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
         super().__init__(tokenizer, tools)
         # Stateful streaming fields
@@ -156,7 +159,25 @@ class Glm4MoeModelToolParser(ToolParser):
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest:
-        """Adjust request parameters for tool call token handling."""
+        """Adjust request parameters for tool call token handling.
+
+        For required/named tool_choice, skip setting structured_outputs
+        because GLM models output tool calls in XML format (per chat
+        template).  Guided decoding would force JSON output, conflicting
+        with the XML format and causing parsing failures.
+        """
+        if request.tools:
+            tc = request.tool_choice
+            if tc == "required" or isinstance(tc, ChatCompletionNamedToolChoiceParam):
+                # Do NOT call super().adjust_request() for required/named,
+                # because it would set structured_outputs and force JSON
+                # output via guided decoding.  GLM models use XML tool-call
+                # syntax (defined in the chat template), so guided decoding
+                # must be skipped to let the model output XML freely.
+                # The tool_parser handles extraction from XML output.
+                if request.tool_choice != "none":
+                    request.skip_special_tokens = False
+                return request
         request = super().adjust_request(request)
         if request.tools and request.tool_choice != "none":
             # Ensure tool call tokens (<tool_call>, </tool_call>) are not skipped

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