fix

Signed-off-by: Roger Wang <hey@rogerw.io>

.buildkite/scripts/run-multi-node-test.sh

@@ -67,7 +67,7 @@ start_nodes() {
         # 3. map the huggingface cache directory to the container
         # 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
         #    starting from 192.168.10.11)
-        docker run -d "$GPU_DEVICES" --shm-size=10.24gb -e HF_TOKEN \
+        docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \
             -v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \
             --network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
             /bin/bash -c "tail -f /dev/null"

.buildkite/test-amd.yaml

@@ -55,9 +55,11 @@ steps:
   grade: Blocking
   source_file_dependencies:
   - vllm/
+  - tests/detokenizer
   - tests/multimodal
   - tests/utils_
   commands:
+  - pytest -v -s detokenizer
   - pytest -v -s -m 'not cpu_test' multimodal
   - pytest -v -s utils_
 
@@ -547,7 +549,7 @@ steps:
   - tests/samplers
   - tests/conftest.py
   commands:
-    - pytest -v -s -m 'not skip_v1' samplers
+    - pytest -v -s -m samplers
 
 - label: LoRA Test %N # 20min each
   timeout_in_minutes: 30
@@ -2213,7 +2215,7 @@ steps:
   - tests/samplers
   - tests/conftest.py
   commands:
-    - pytest -v -s -m 'not skip_v1' samplers
+    - pytest -v -s -m samplers
 
 - label: LoRA Test %N # 20min each
   timeout_in_minutes: 30

.buildkite/test-pipeline.yaml

@@ -51,9 +51,11 @@ steps:
   mirror_hardwares: [amdexperimental]
   source_file_dependencies:
   - vllm/
+  - tests/detokenizer
   - tests/multimodal
   - tests/utils_
   commands:
+  - pytest -v -s detokenizer
   - pytest -v -s -m 'not cpu_test' multimodal
   - pytest -v -s utils_
 

.buildkite/test_areas/distributed.yaml

@@ -165,6 +165,7 @@ steps:
   num_devices: 2
   num_nodes: 2
   no_plugin: true
+  optional: true # TODO: revert once infra issue solved
   source_file_dependencies:
   - vllm/distributed/
   - vllm/engine/

.buildkite/test_areas/misc.yaml

@@ -108,9 +108,11 @@ steps:
   timeout_in_minutes: 50
   source_file_dependencies:
   - vllm/
+  - tests/detokenizer
   - tests/multimodal
   - tests/utils_
   commands:
+  - pytest -v -s detokenizer
   - pytest -v -s -m 'not cpu_test' multimodal
   - pytest -v -s utils_
 

.buildkite/test_areas/samplers.yaml

@@ -18,4 +18,4 @@ steps:
       depends_on:
       - image-build-amd
       commands:
-      - pytest -v -s -m 'not skip_v1' samplers
+      - pytest -v -s -m samplers

CMakeLists.txt

@@ -771,6 +771,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
     endif()
   endif()
 
+  # DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
+    set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${DSV3_FUSED_A_GEMM_SRC}"
+      CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
+    list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
+    list(APPEND VLLM_GPU_FLAGS "-DENABLE_DSV3_FUSED_A_GEMM=1")
+    message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
+  else()
+    message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
+                   "in CUDA target architectures.")
+  endif()
+
   # moe_data.cu is used by all CUTLASS MoE kernels.
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")

csrc/cache_kernels.cu

@@ -1305,7 +1305,8 @@ void indexer_k_quant_and_cache(
   const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
-  DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
+  static const std::string kv_cache_dtype = "fp8_e4m3";
+  DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype,
                              CALL_INDEXER_K_QUANT_AND_CACHE);
 }
 

csrc/dsv3_fused_a_gemm.cu

@@ -0,0 +1,747 @@
+/*
+ * Adapted from
+ * https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
+ * which was adapted from
+ * https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu
+ *
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2021, NAVER Corp.  Authored by CLOVA.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <torch/all.h>
+
+#include "core/registration.h"
+
+#include <cstdlib>
+#include <mutex>
+
+namespace {
+
+inline int getSMVersion() {
+  auto* props = at::cuda::getCurrentDeviceProperties();
+  return props->major * 10 + props->minor;
+}
+
+inline bool getEnvEnablePDL() {
+  static std::once_flag flag;
+  static bool enablePDL = false;
+  std::call_once(flag, [&]() {
+    if (getSMVersion() >= 90) {
+      char const* env = std::getenv("TRTLLM_ENABLE_PDL");
+      enablePDL = env && env[0] == '1' && env[1] == '\0';
+    }
+  });
+  return enablePDL;
+}
+
+}  // namespace
+
+using bf16_t = __nv_bfloat16;
+
+__device__ void hmma_16_8_16_f32acc_bf16ab(float (&d_reg)[4],
+                                           const bf16_t (&a_reg)[8],
+                                           const bf16_t (&b_reg)[4],
+                                           float const (&c_reg)[4]) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t a0 = *reinterpret_cast<uint32_t const*>(a_reg + 0);
+  uint32_t a1 = *reinterpret_cast<uint32_t const*>(a_reg + 2);
+  uint32_t a2 = *reinterpret_cast<uint32_t const*>(a_reg + 4);
+  uint32_t a3 = *reinterpret_cast<uint32_t const*>(a_reg + 6);
+  uint32_t b0 = *reinterpret_cast<uint32_t const*>(b_reg + 0);
+  uint32_t b1 = *reinterpret_cast<uint32_t const*>(b_reg + 2);
+  asm volatile(
+      "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+      "{%0,  %1,  %2,  %3},"
+      "{%4,  %5,  %6,  %7},"
+      "{%8,  %9},"
+      "{%10, %11, %12, %13};\n"
+      : "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3])
+      : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(d_reg[0]),
+        "f"(d_reg[1]), "f"(d_reg[2]), "f"(d_reg[3]));
+#endif
+}
+
+extern "C" {
+__device__ uint32_t __nvvm_get_smem_pointer(void*);
+}
+
+__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  if (pred) {
+    uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr);
+    asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(
+                     smemPtrAsUint32),
+                 "l"(gPtr), "n"(16));
+  }
+#endif
+}
+
+__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  asm volatile(
+      "ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+      : "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3])
+      : "r"(__nvvm_get_smem_pointer(smem_ptr)));
+#endif
+}
+
+template <class Type>
+__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) {
+  uint32_t row_idx = *reinterpret_cast<uint32_t*>(&row_idx_);
+  uint32_t col_idx = *reinterpret_cast<uint32_t*>(&col_idx_);
+  row_idx = row_idx % 8;
+  row_idx = row_idx * (16 / sizeof(Type));
+  col_idx = col_idx ^ row_idx;
+  return *reinterpret_cast<int*>(&col_idx);
+}
+
+__device__ void initialize_barrier(
+    uint64_t* smem_barrier,  // 64 bits user-manged barrier in smem
+    int thread_count =
+        1)  // Thread count expected to arrive/wait on this barrier
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr),
+               "r"(thread_count));
+#endif
+}
+
+// Barrier wait
+__device__ void wait_barrier(
+    uint64_t* smem_barrier,  // 64 bits user-manged barrier in smem
+    int phase_bit)           // Current phase bit the barrier waiting to flip
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile(
+      "{\n"
+      ".reg .pred                P1;\n"
+      "LAB_WAIT:\n"
+      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
+      "@P1                       bra DONE;\n"
+      "bra                   LAB_WAIT;\n"
+      "DONE:\n"
+      "}\n" ::"r"(smem_int_ptr),
+      "r"(phase_bit));
+#endif
+}
+
+__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t wait_complete;
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr);
+  asm volatile(
+      "{\n\t"
+      ".reg .pred P1; \n\t"
+      "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
+      "selp.b32 %0, 1, 0, P1; \n\t"
+      "}"
+      : "=r"(wait_complete)
+      : "r"(smem_int_ptr), "r"(phase_bit));
+  return static_cast<bool>(wait_complete);
+#endif
+  return false;
+}
+
+// Barrier arrive
+__device__ void arrive_barrier(
+    uint64_t* smem_barrier)  // 64 bits user-manged barrier in smem
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile(
+      "{\n"
+      ".reg .b64 state; \n"
+      "mbarrier.arrive.shared::cta.b64   state, [%0];\n"
+      "}\n" ::"r"(smem_int_ptr));
+#endif
+}
+
+__device__ void ldgsts_arrive(uint64_t* smem_barrier) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
+  asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];"
+               :
+               : "r"(smem_int_ptr));
+#endif
+}
+
+template <int gemm_k, int tile_m, int tile_k, int stage_cnt>
+struct GmemLoaderA {
+  static constexpr int elem_bytes = 2;
+  static constexpr int vec_bytes = 16;
+  static constexpr int vec_elems = vec_bytes / elem_bytes;
+  static constexpr int thread_cnt = 64;
+  static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0);
+  static constexpr int a_inst_cnt_per_iter =
+      (tile_m * tile_k) / (vec_elems * thread_cnt);
+  static_assert(gemm_k % tile_k == 0);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+
+  // Extra params to keep the order of k reduction...
+  static constexpr int mma_warp_cnt = 4;
+  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
+  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
+
+ private:
+  __device__ int k_project(int tile_k_idx) {
+    return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
+           (tile_k_idx % per_mma_warp_k);
+  }
+
+ public:
+  __device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_,
+                         uint64_t* smem_barrier_)
+      : gmem_a(gmem_a_local_),
+        smem_a(smem_a_),
+        smem_barrier(smem_barrier_),
+        local_tid(threadIdx.x % thread_cnt) {}
+
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  // swizzle, that's what we want.
+  #pragma unroll
+    for (int i = 0; i < a_inst_cnt_per_iter; i++) {
+      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+      int m_idx = linear_idx / tile_k;
+      int k_idx = linear_idx % tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
+      a_smem_offsets[i] = m_idx * tile_k + k_idx;
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  #pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      if (need_wait) {
+        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
+      }
+      int next_stage_idx = stage_idx + 1;
+      int next_phase_bit =
+          next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
+      if (loop_idx != k_iter_cnt - 1) {
+        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
+                                      next_phase_bit);
+      }
+
+  #pragma unroll
+      for (int i = 0; i < a_inst_cnt_per_iter; i++) {
+        int smem_offset = a_smem_offsets[i];
+        bf16_t* smem_ptr_this_iter =
+            smem_a + stage_idx * tile_m * tile_k + smem_offset;
+        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+        int m_idx = linear_idx / tile_k;
+        int k_idx = linear_idx % tile_k;
+        int gmem_offset = m_idx * gemm_k + k_project(k_idx);
+        bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset;
+        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true);
+      }
+      ldgsts_arrive(smem_barrier + stage_idx * 2);
+
+      stage_idx = next_stage_idx;
+      phase_bit = next_phase_bit;
+      gmem_a += per_mma_warp_k;
+    }
+#endif
+  }
+
+  bf16_t const* gmem_a;
+  bf16_t* smem_a;
+  uint64_t* smem_barrier;
+  int local_tid;
+  int stage_idx = 0;
+  int phase_bit = 1;
+  bool need_wait = true;
+
+  // per smem_stage, store with swizzle information
+  int a_smem_offsets[a_inst_cnt_per_iter];
+};
+
+template <int gemm_k, int tile_n, int tile_k, int stage_cnt>
+struct GmemLoaderB {
+  static constexpr int elem_bytes = 2;
+  static constexpr int vec_bytes = 16;
+  static constexpr int vec_elems = vec_bytes / elem_bytes;
+  static constexpr int thread_cnt = 64;
+  static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0);
+  static constexpr int b_inst_cnt_per_iter =
+      (tile_n * tile_k) / (vec_elems * thread_cnt);
+  static_assert(gemm_k % tile_k == 0);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+
+  // Extra params to keep the order of k reduction...
+  static constexpr int mma_warp_cnt = 4;
+  static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
+  static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
+
+ private:
+  __device__ int k_project(int tile_k_idx) {
+    return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
+           (tile_k_idx % per_mma_warp_k);
+  }
+
+ public:
+  __device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_,
+                         uint64_t* smem_barrier_, int gemm_n_)
+      : gmem_b(gmem_b_local_),
+        smem_b(smem_b_),
+        smem_barrier(smem_barrier_),
+        gemm_n(gemm_n_),
+        local_tid(threadIdx.x % thread_cnt) {}
+
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  // swizzle, that's what we want.
+  #pragma unroll
+    for (int i = 0; i < b_inst_cnt_per_iter; i++) {
+      int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+      int n_idx = linear_idx / tile_k;
+      int k_idx = linear_idx % tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
+      b_smem_offsets[i] = n_idx * tile_k + k_idx;
+      preds[i] = n_idx < gemm_n;
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    asm volatile("griddepcontrol.wait;");
+  #pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      if (need_wait) {
+        wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
+      }
+      int next_stage_idx = stage_idx + 1;
+      int next_phase_bit =
+          next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
+      if (loop_idx != k_iter_cnt - 1) {
+        need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
+                                      next_phase_bit);
+      }
+  #pragma unroll
+      for (int i = 0; i < b_inst_cnt_per_iter; i++) {
+        int smem_offset = b_smem_offsets[i];
+        bf16_t* smem_ptr_this_iter =
+            smem_b + stage_idx * tile_n * tile_k + smem_offset;
+        int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
+        int n_idx = linear_idx / tile_k;
+        int k_idx = linear_idx % tile_k;
+        int gmem_offset = n_idx * gemm_k + k_project(k_idx);
+        bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
+        ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
+      }
+      ldgsts_arrive(smem_barrier + stage_idx * 2);
+
+      stage_idx = next_stage_idx;
+      phase_bit = next_phase_bit;
+      gmem_b += per_mma_warp_k;
+    }
+#endif
+  }
+
+  bf16_t const* gmem_b;
+  bf16_t* smem_b;
+  uint64_t* smem_barrier;
+  int gemm_n;
+  int local_tid;
+  int stage_idx = 0;
+  int phase_bit = 1;
+  bool need_wait = true;
+
+  // per smem_stage, store with swizzle information
+  int b_smem_offsets[b_inst_cnt_per_iter];
+  uint32_t preds[b_inst_cnt_per_iter];
+};
+
+template <int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k,
+          int stage_cnt>
+struct MmaComputer {
+  static constexpr int elem_bytes = 2;
+  static constexpr int thread_cnt = 128;
+  static_assert(gemm_k % tile_k == 0);
+  static_assert(tile_k % (thread_cnt / 32) == 0);
+  static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32);
+  static constexpr int k_iter_cnt = gemm_k / tile_k;
+  static constexpr int k_phase_cnt = per_warp_tile_k / 16;
+  static constexpr int m_iter_cnt = (tile_m + 15) / 16;
+  static constexpr int n_iter_cnt =
+      (tile_n + 7) /
+      8;  // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
+  static_assert(m_iter_cnt == 1);
+  static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
+
+  __device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
+                         bf16_t* smem_b_, uint64_t* smem_barrier_,
+                         int warp_idx_, int gemm_n_)
+      : gmem_c(gmem_c_local_),
+        smem_a(smem_a_),
+        smem_b(smem_b_),
+        smem_barrier(smem_barrier_),
+        warp_idx(warp_idx_ - (thread_cnt / 32)),
+        gemm_n(gemm_n_) {}
+
+ private:
+  __device__ constexpr int internal_b_atom_func(int tid) {
+    if constexpr (tile_n < 8) {
+      return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n;
+    } else {
+      return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8));
+    }
+  }
+
+ public:
+  __device__ void prepare() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  #pragma unroll
+    for (int i = 0; i < k_phase_cnt; i++) {
+      int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
+      int m_idx = linear_idx % tile_m;
+      int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
+      k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
+      a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
+    }
+  #pragma unroll
+    for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+  #pragma unroll
+      for (int i = 0; i < k_phase_cnt; i += 2) {  // Special i+=2 for B.
+        int linear_idx =
+            internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8;
+        int n_idx = linear_idx % tile_n;
+        int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k;
+        k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
+        b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx;
+      }
+    }
+#endif
+  }
+
+  __device__ void issue_mainloop() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  #pragma unroll 1
+    for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
+      wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
+
+  #pragma unroll
+      for (int i = 0; i < k_phase_cnt; i++) {
+        int smem_offset = a_smem_offsets[0][i];
+        bf16_t* smem_ptr_this_iter =
+            smem_a + stage_idx * tile_m * tile_k + smem_offset;
+        ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
+      }
+
+  #pragma unroll
+      for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+  #pragma unroll
+        for (int i = 0; i < k_phase_cnt; i += 2) {
+          int smem_offset = b_smem_offsets[n_iter_idx][i];
+          bf16_t* smem_ptr_this_iter =
+              smem_b + stage_idx * tile_n * tile_k + smem_offset;
+          ldsm_x4(smem_ptr_this_iter,
+                  reinterpret_cast<uint32_t*>(b_reg[n_iter_idx][i]));
+        }
+      }
+
+  #pragma unroll
+      for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
+  #pragma unroll
+        for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
+          hmma_16_8_16_f32acc_bf16ab(
+              acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
+              b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
+        }
+      }
+      ::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
+      stage_idx += 1;
+      phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
+      stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
+    }
+#endif
+  }
+
+  __device__ void epi() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
+    // reorganize the acc_reg
+    constexpr int thread_m = 2;
+    constexpr int thread_n = 2 * n_iter_cnt;
+    constexpr int cta_mma_n = n_iter_cnt * 8;
+    float acc_reg_reorg[thread_m][thread_n];
+
+    for (int i = 0; i < thread_m; i++) {
+      for (int j = 0; j < thread_n; j++) {
+        acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
+      }
+    }
+
+    // 4 x cosize(smem_c_layout)
+    float* smem_c = reinterpret_cast<float*>(smem_a);
+    // coord -> index
+    auto smem_c_index_func = [&](int m_idx, int n_idx) {
+      int group_rows = 32 / cta_mma_n;
+      int group_cnt = 2;
+      return (m_idx % group_rows * cta_mma_n) +
+             (m_idx / group_rows * (32 + group_cnt)) + n_idx;
+    };
+    constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2);
+
+  // This should be optimized to STS.64 but can not be STS.128 due to the bank
+  // index.
+  #pragma unroll
+    for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
+  #pragma unroll
+      for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
+        int m_idx = (lane_idx / 4) + m_idx_thread * 8;
+        int n_idx =
+            ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
+        smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
+            acc_reg_reorg[m_idx_thread][n_idx_thread];
+      }
+    }
+    asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
+
+    if (warp_idx == 0) {
+      constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32;
+      float acc_final[final_acc_reg_cnt]{};
+
+  #pragma unroll
+      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
+        int linear_idx = reg_idx * 32 + lane_idx;
+        int m_idx = linear_idx % tile_m;
+        int n_idx = linear_idx / tile_m;
+        acc_final[reg_idx] +=
+            smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] +
+            smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] +
+            smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] +
+            smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c];
+      }
+
+  #pragma unroll
+      for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
+        int linear_idx = reg_idx * 32 + lane_idx;
+        int m_idx = linear_idx % tile_m;
+        int n_idx = linear_idx / tile_m;
+        if (m_idx < tile_m && n_idx < gemm_n) {
+          gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx];
+        }
+      }
+    }
+#endif
+  }
+
+  bf16_t* gmem_c;
+  bf16_t* smem_a;
+  bf16_t* smem_b;
+  uint64_t* smem_barrier;
+  int warp_idx;
+  int gemm_n;
+  int stage_idx = 0;
+  int phase_bit = 0;
+  int lane_idx = threadIdx.x % 32;
+  int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k;
+
+  int a_smem_offsets[m_iter_cnt][k_phase_cnt];
+  int b_smem_offsets[n_iter_cnt][k_phase_cnt];
+
+  bf16_t a_reg[m_iter_cnt][k_phase_cnt][8];
+  bf16_t b_reg[n_iter_cnt][k_phase_cnt][4];
+  float acc_reg[m_iter_cnt][n_iter_cnt][4]{};
+};
+
+// AB swapped, kernel is k-major, k-major, m-major
+template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n,
+          int tile_k, int stage_cnt>
+__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
+    bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  constexpr int load_thread_cnt = 128;
+  constexpr int compute_thread_cnt = 128;
+  constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt;
+  (void)thread_cnt;
+  static_assert(gemm_m % 16 == 0);
+  static_assert(gemm_k % tile_k == 0);
+  static_assert(gemm_m % tile_m == 0);
+  static_assert(
+      tile_k == 128 || tile_k == 256 || tile_k == 512 ||
+      tile_k == 1024);  // tile_k must be larger than 64 since 4 warp splitK.
+  static_assert(tile_m == 16);
+  constexpr int g2s_vec_bytes = 16;
+  constexpr int a_elem_bytes = 2;
+  constexpr int b_elem_bytes = 2;
+  static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k *
+                    stage_cnt <=
+                225 * 1024);
+  static_assert((tile_m * tile_k * a_elem_bytes) %
+                    (load_thread_cnt * g2s_vec_bytes) ==
+                0);
+  static_assert((tile_n * tile_k * b_elem_bytes) %
+                    (load_thread_cnt * g2s_vec_bytes) ==
+                0);
+
+  extern __shared__ char smem[];
+  uint64_t* smem_barrier = reinterpret_cast<uint64_t*>(
+      smem);  // producer,consumer; producer,consumer; ...
+  bf16_t* smem_a = reinterpret_cast<bf16_t*>(smem + (stage_cnt * 8 * 2 + 1024) /
+                                                        1024 * 1024);
+  bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt;
+
+  int cta_m_idx = tile_m * blockIdx.x;
+  int cta_n_idx = tile_n * blockIdx.y;
+  bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
+  bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
+  bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
+
+  int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+
+  if (warp_idx == 4) {
+    for (int i = 0; i < stage_cnt; i++) {
+      initialize_barrier(smem_barrier + i * 2 + 0,
+                         load_thread_cnt);  // producer
+      initialize_barrier(smem_barrier + i * 2 + 1,
+                         compute_thread_cnt);  // consumer
+    }
+  }
+  __syncthreads();
+
+  if (warp_idx < 2) {
+    GmemLoaderA<gemm_k, tile_m, tile_k, stage_cnt> a_loader(
+        gmem_a_local, smem_a, smem_barrier);
+    a_loader.prepare();
+    a_loader.issue_mainloop();
+  } else if (warp_idx < 4) {
+    GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(
+        gmem_b_local, smem_b, smem_barrier, gemm_n);
+    b_loader.prepare();
+    b_loader.issue_mainloop();
+  } else {
+    MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
+        gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
+    mma_computer.prepare();
+    mma_computer.issue_mainloop();
+    mma_computer.epi();
+  }
+  asm volatile("griddepcontrol.launch_dependents;");
+#endif
+}
+
+template <typename T, int kHdIn, int kHdOut, int kTileN>
+void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
+                      cudaStream_t const stream) {
+  constexpr int gemm_m = kHdOut;  // 2112
+  int const gemm_n = num_tokens;  // 1-16
+  constexpr int gemm_k = kHdIn;   // 7168
+  constexpr int batch_size = 1;
+  std::swap(mat_a, mat_b);
+  constexpr int tile_m = 16;
+  constexpr int tile_n = kTileN;                        // 8 or 16
+  constexpr int tile_k = std::max(256, 1024 / tile_n);  // 256
+  constexpr int max_stage_cnt =
+      1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
+  constexpr int k_iter_cnt = gemm_k / tile_k;
+  constexpr int stage_cnt =
+      k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt;
+  int cta_m_cnt = gemm_m / tile_m;
+  int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n;
+  constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024;
+  constexpr int smem_bytes =
+      ((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) /
+      1024 * 1024;
+
+  dim3 grid(cta_m_cnt, cta_n_cnt, 1);
+  dim3 block_size(256);
+  cudaLaunchConfig_t config;
+  config.gridDim = grid;
+  config.blockDim = block_size;
+  config.dynamicSmemBytes = smem_bytes;
+  config.stream = stream;
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
+  config.numAttrs = 1;
+  config.attrs = attrs;
+  if (smem_bytes >= (48 * 1024)) {
+    cudaFuncSetAttribute(fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
+                                             tile_n, tile_k, stage_cnt>,
+                         cudaFuncAttributeMaxDynamicSharedMemorySize,
+                         smem_bytes);
+  }
+  cudaLaunchKernelEx(&config,
+                     fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
+                                         tile_n, tile_k, stage_cnt>,
+                     output, mat_a, mat_b, gemm_n);
+}
+
+template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
+    cudaStream_t);
+
+template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
+    __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
+    cudaStream_t);
+
+void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
+                       torch::Tensor const& mat_b) {
+  TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
+  int const num_tokens = mat_a.size(0);
+  int const hd_in = mat_a.size(1);
+  int const hd_out = mat_b.size(1);
+
+  constexpr int kHdIn = 7168;
+  constexpr int kHdOut = 2112;
+  TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
+              "required 1 <= mat_a.shape[0] <= 16")
+  TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
+  TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
+  TORCH_CHECK(output.size(0) == num_tokens,
+              "required output.shape[0] == mat_a.shape[0]")
+  TORCH_CHECK(output.size(1) == hd_out,
+              "required output.shape[1] == mat_b.shape[1]")
+
+  TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
+  TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
+  TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
+
+  TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
+                  mat_b.scalar_type() == torch::kBFloat16,
+              "Only BFloat16 input dtype is supported")
+  TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
+              "Only BFloat16 output dtype is supported")
+
+  TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
+
+  auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
+  if (num_tokens <= 8) {
+    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
+        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
+        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
+        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
+        stream);
+  } else {
+    invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
+        reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
+        reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
+        reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
+        stream);
+  }
+}

csrc/ops.h

@@ -315,7 +315,9 @@ void silu_and_mul_scaled_fp4_experts_quant(
 void per_token_group_quant_fp8(const torch::Tensor& input,
                                torch::Tensor& output_q, torch::Tensor& output_s,
                                int64_t group_size, double eps, double fp8_min,
-                               double fp8_max, bool scale_ue8m0);
+                               double fp8_max, bool scale_ue8m0,
+                               bool dummy_is_scale_transposed,
+                               bool dummy_is_tma_aligned);
 
 void per_token_group_quant_int8(const torch::Tensor& input,
                                 torch::Tensor& output_q,
@@ -408,3 +410,8 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
                    int64_t quant_level, bool cast_bf2half = false);
 int64_t qr_max_size();
 #endif
+
+#ifndef USE_ROCM
+void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
+                       torch::Tensor const& mat_b);
+#endif

csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

@@ -97,7 +97,7 @@ __global__ void rms_norm_per_block_quant_kernel(
     scalar_t const* __restrict__ input,   // [..., hidden_size]
     scalar_t const* __restrict__ weight,  // [hidden_size]
     float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
-    scalar_t* __restrict__ residual = nullptr) {
+    scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
   float rms;
   // Compute RMS
   // Always able to vectorize due to constraints on hidden_size
@@ -108,7 +108,8 @@ __global__ void rms_norm_per_block_quant_kernel(
   // Always able to vectorize due to constraints on hidden_size and group_size
   vllm::vectorized::compute_dynamic_per_token_scales<
       scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(
-      nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual);
+      nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual,
+      outer_scale_stride);
 
   // RMS Norm + Quant
   // Always able to vectorize due to constraints on hidden_size
@@ -119,7 +120,8 @@ __global__ void rms_norm_per_block_quant_kernel(
   vllm::vectorized::norm_and_quant<
       scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,
       has_residual, is_scale_transposed, group_size>(
-      out, input, weight, rms, scales, hidden_size, residual);
+      out, input, weight, rms, scales, hidden_size, residual,
+      outer_scale_stride);
 }
 
 }  // namespace vllm
@@ -225,7 +227,8 @@ void rms_norm_per_block_quant_dispatch(
                                                  : nullptr,
                             var_epsilon, hidden_size,
                             has_residual ? residual->data_ptr<scalar_in_t>()
-                                         : nullptr);
+                                         : nullptr,
+                            scales.stride(1));
                   });
             });
           });
@@ -257,6 +260,11 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
   TORCH_CHECK(group_size == 128 || group_size == 64,
               "Unsupported group size: ", group_size);
 
+  if (scales.stride(1) > 1) {
+    TORCH_CHECK(is_scale_transposed,
+                "Outer scale stride must be 1 when scales are not transposed");
+  }
+
   rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
                                     var_epsilon, scale_ub, residual,
                                     is_scale_transposed);

csrc/quantization/fused_kernels/layernorm_utils.cuh

@@ -74,7 +74,7 @@ __device__ void compute_dynamic_per_token_scales(
     scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
     float const rms, float const* __restrict__ scale_ub,
     int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
-    int32_t const group_size = 0) {
+    int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
   float block_absmax_val_maybe = 0.0f;
   constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
   __syncthreads();
@@ -133,7 +133,9 @@ __device__ void compute_dynamic_per_token_scales(
       scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
       // Global output store
       if constexpr (is_scale_transposed) {
-        all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
                          blockIdx.x] = scale;
       } else {
         all_token_scales[blockIdx.x * num_groups +
@@ -180,13 +182,11 @@ __device__ void compute_dynamic_per_token_scales(
 
 template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
           bool has_residual = false, bool is_scale_transposed = false>
-__device__ void norm_and_quant(scalar_out_t* __restrict__ output,
-                               scalar_t const* __restrict__ input,
-                               scalar_t const* __restrict__ weight,
-                               float const rms, float* const scale,
-                               int32_t const hidden_size,
-                               scalar_t* __restrict__ residual = nullptr,
-                               int32_t const group_size = 0) {
+__device__ void norm_and_quant(
+    scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
+    scalar_t const* __restrict__ weight, float const rms, float* const scale,
+    int32_t const hidden_size, scalar_t* __restrict__ residual = nullptr,
+    int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
   int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
 
   for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
@@ -202,7 +202,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
     int64_t scale_idx = 0;
     if (group_size > 0) {
       if constexpr (is_scale_transposed) {
-        scale_idx = (i / group_size) * gridDim.x + blockIdx.x;
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        scale_idx = (i / group_size) * scale_rows + blockIdx.x;
       } else {
         scale_idx = blockIdx.x * (hidden_size / group_size) + i / group_size;
       }
@@ -286,8 +288,8 @@ __device__ void compute_dynamic_per_token_scales(
     float* __restrict__ token_scale, float* __restrict__ all_token_scales,
     scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
     float const rms, float const* __restrict__ scale_ub,
-    int32_t const hidden_size,
-    scalar_t const* __restrict__ residual = nullptr) {
+    int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
+    int64_t outer_scale_stride = 1) {
   constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
 
   const int VEC_SIZE = 4;
@@ -382,7 +384,9 @@ __device__ void compute_dynamic_per_token_scales(
       scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
       // Global output store
       if constexpr (is_scale_transposed) {
-        all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
                          blockIdx.x] = scale;
       } else {
         all_token_scales[blockIdx.x * num_groups +
@@ -463,7 +467,8 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
                                scalar_t const* __restrict__ weight,
                                float const rms, float* const scale,
                                int32_t const hidden_size,
-                               scalar_t* __restrict__ residual = nullptr) {
+                               scalar_t* __restrict__ residual = nullptr,
+                               int64_t outer_scale_stride = 1) {
   int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
 
   // Vectorized input/output/weight/residual to better utilize memory bandwidth.
@@ -516,7 +521,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
       int64_t const num_groups = hidden_size / group_size;
       int64_t scale_idx = 0;
       if constexpr (is_scale_transposed) {
-        scale_idx = (i * VEC_SIZE / group_size) * gridDim.x + blockIdx.x;
+        int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
+                                   outer_scale_stride * outer_scale_stride;
+        scale_idx = (i * VEC_SIZE / group_size) * scale_rows + blockIdx.x;
       } else {
         scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size;
       }

csrc/quantization/w8a8/fp8/per_token_group_quant.cu

@@ -379,7 +379,9 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
 void per_token_group_quant_fp8(const torch::Tensor& input,
                                torch::Tensor& output_q, torch::Tensor& output_s,
                                int64_t group_size, double eps, double fp8_min,
-                               double fp8_max, bool scale_ue8m0) {
+                               double fp8_max, bool scale_ue8m0,
+                               bool dummy_is_scale_transposed = false,
+                               bool dummy_is_tma_aligned = false) {
   per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
                              fp8_min, fp8_max, scale_ue8m0);
 }

csrc/torch_bindings.cpp

@@ -239,6 +239,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
   // Quantization ops
 #ifndef USE_ROCM
+  // DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
+  ops.def(
+      "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
+  ops.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
+
   // Quantized GEMM for AWQ.
   ops.def(
       "awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
@@ -643,11 +648,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
 #ifndef USE_ROCM
   // Compute per-token-group FP8 quantized tensor and scaling factor.
+  // The dummy arguments are here so we can correctly fuse with RMSNorm.
   ops.def(
       "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! "
       "output_s, "
       "int group_size, float eps, float fp8_min, float fp8_max, bool "
-      "scale_ue8m0) -> ()");
+      "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned "
+      ") -> ()");
   ops.impl("per_token_group_fp8_quant", torch::kCUDA,
            &per_token_group_quant_fp8);
 

docs/design/plugin_system.md

@@ -155,3 +155,4 @@ The interface for the model/module may change during vLLM's development. If you
     - `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It has been removed in v0.13.0.
     - `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.v1.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
     - `seed_everything` platform interface is deprecated. It has been removed in v0.16.0. Please use `vllm.utils.torch_utils.set_random_seed` instead.
+    - `prompt` in `Platform.validate_request` is deprecated and will be removed in v0.18.0.

pyproject.toml

@@ -113,7 +113,6 @@ markers = [
     "cpu_test: mark test as CPU-only test",
     "split: run this test as part of a split",
     "distributed: run this test only in distributed GPU tests",
-    "skip_v1: do not run this test with v1",
     "optional: optional tests that are automatically skipped, include --optional to run them",
 ]
 

requirements/kv_connectors.txt

@@ -1,2 +1,3 @@
 lmcache >= 0.3.9
 nixl >= 0.7.1 # Required for disaggregated prefill
+mooncake-transfer-engine >= 0.3.8

tests/compile/fusions_e2e/test_tp1_quant.py

@@ -50,10 +50,9 @@ def test_tp1_fp8_fusions(
     run_e2e_fusion_test,
     monkeypatch,
 ):
-    if use_deepgemm:
-        # TODO(luka/eliza) DeepGEMM uses different quants, matching not supported
+    if use_deepgemm and is_blackwell():
+        # TODO(luka) DeepGEMM uses different quants, matching not supported
         #  - on Blackwell, uses a special quant fp8, currently not supported
-        #  - on Hopper, tma-aligned scales inhibit matching (fix WIP)
         pytest.skip("DeepGEMM & quant matching not currently supported")
 
     matches = matches_fn(n_layers)
@@ -66,7 +65,6 @@ def test_tp1_fp8_fusions(
     model_kwargs["hf_overrides"] = hf_overrides(n_layers)
     model_kwargs["load_format"] = "dummy"
     model_kwargs["max_model_len"] = 1024
-
     compilation_config = dict(
         use_inductor_graph_partition=inductor_graph_partition,
         custom_ops=custom_ops.split(","),

tests/detokenizer/test_disable_detokenization.py

@@ -7,7 +7,6 @@ from vllm.entrypoints.llm import LLM
 from vllm.sampling_params import SamplingParams
 
 
-@pytest.mark.skip_v1
 @pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
 def test_computed_prefix_blocks(model: str):
     # This test checks if the engine generates completions both with and

tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml

@@ -0,0 +1,8 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
+accuracy_threshold: 0.29
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
+env:
+  VLLM_USE_FLASHINFER_MOE_FP8: "1"
+  VLLM_FLASHINFER_MOE_BACKEND: "latency"

tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml

@@ -0,0 +1,8 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"
+accuracy_threshold: 0.29
+num_questions: 1319
+num_fewshot: 5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
+env:
+  VLLM_USE_FLASHINFER_MOE_FP4: "1"
+  VLLM_FLASHINFER_MOE_BACKEND: "throughput"

tests/evals/gsm8k/configs/moe-refactor/config-b200.txt

@@ -13,3 +13,5 @@ Llama-4-Scout-BF16-fi-cutlass.yaml
 Llama-4-Scout-BF16-triton.yaml
 Mixtral-8x7B-BF16-fi-cutlass.yaml
 Mixtral-8x7B-BF16-triton.yaml
+Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml
+Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml

tests/kernels/attention/test_triton_unified_attention.py

@@ -10,7 +10,7 @@ from vllm.utils.math_utils import next_power_of_2
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.ops.triton_unified_attention import unified_attention
 
-NUM_HEADS = [(4, 4), (8, 2)]
+NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
 HEAD_SIZES = [128, 256]
 BLOCK_SIZES = [16]
 
@@ -20,6 +20,8 @@ QDTYPES = (
     if not current_platform.is_rocm()
     else [None, torch.float8_e4m3fnuz]
 )
+FP8_DTYPE = current_platform.fp8_dtype()
+
 # one value large enough to test overflow in index calculation.
 # one value small enough to test the schema op check
 NUM_BLOCKS = [32768, 2048]
@@ -217,3 +219,127 @@ def test_triton_unified_attn(
         torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
         f"{torch.max(torch.abs(output - ref_output))}",
     )
+
+
+@pytest.mark.parametrize(
+    "seq_lens",
+    [
+        [(1, 1328), (5, 18), (129, 463)],
+        [(1, 523), (1, 37), (1, 2011)],
+        [(1, 1)] * 533,
+        [(533, 533)] * 533,
+    ],
+)
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
+@pytest.mark.parametrize("head_size", HEAD_SIZES)
+@pytest.mark.parametrize("block_size", BLOCK_SIZES)
+@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
+@pytest.mark.parametrize("soft_cap", [None, 50.0])
+@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
+@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
+@torch.inference_mode()
+def test_triton_unified_attn_fp16_input_fp8_output(
+    seq_lens: list[tuple[int, int]],
+    num_heads: tuple[int, int],
+    head_size: int,
+    sliding_window: int | None,
+    block_size: int,
+    soft_cap: float | None,
+    num_blocks: int,
+    seq_threshold_3D: int,
+) -> None:
+    """Test with fp16 input and fp8 output using output_scale."""
+    torch.set_default_device("cuda")
+
+    set_random_seed(0)
+    num_seqs = len(seq_lens)
+    query_lens = [x[0] for x in seq_lens]
+    kv_lens = [x[1] for x in seq_lens]
+    num_query_heads = num_heads[0]
+    num_kv_heads = num_heads[1]
+    assert num_query_heads % num_kv_heads == 0
+    max_query_len = max(query_lens)
+    max_kv_len = max(kv_lens)
+    window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
+    scale = head_size**-0.5
+
+    dtype = torch.float16
+    query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
+    key_cache = torch.randn(
+        num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
+    )
+    value_cache = torch.randn_like(key_cache)
+    cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
+        dim=0, dtype=torch.int32
+    )
+    kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
+
+    max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
+    block_tables = torch.randint(
+        0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
+    )
+
+    output = torch.empty(sum(query_lens), num_query_heads, head_size, dtype=FP8_DTYPE)
+
+    output_scale = torch.tensor(0.5, dtype=torch.float32)
+
+    num_par_softmax_segments = 16
+    head_size_padded = next_power_of_2(head_size)
+    softmax_segm_output = torch.empty(
+        (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
+        dtype=torch.float32,
+    )
+    softmax_segm_max = torch.empty(
+        (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
+        dtype=torch.float32,
+    )
+    softmax_segm_expsum = torch.empty(
+        (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
+        dtype=torch.float32,
+    )
+
+    unified_attention(
+        q=query,
+        k=key_cache,
+        v=value_cache,
+        out=output,
+        cu_seqlens_q=cu_query_lens,
+        seqused_k=kv_lens_tensor,
+        max_seqlen_q=max_query_len,
+        max_seqlen_k=max_kv_len,
+        softmax_scale=scale,
+        causal=True,
+        window_size=window_size,
+        block_table=block_tables,
+        softcap=soft_cap if soft_cap is not None else 0,
+        q_descale=None,
+        k_descale=None,
+        v_descale=None,
+        output_scale=output_scale,
+        seq_threshold_3D=seq_threshold_3D,
+        num_par_softmax_segments=num_par_softmax_segments,
+        softmax_segm_output=softmax_segm_output,
+        softmax_segm_max=softmax_segm_max,
+        softmax_segm_expsum=softmax_segm_expsum,
+    )
+
+    ref_output = ref_paged_attn(
+        query=query,
+        key_cache=key_cache,
+        value_cache=value_cache,
+        query_lens=query_lens,
+        kv_lens=kv_lens,
+        block_tables=block_tables,
+        scale=scale,
+        sliding_window=sliding_window,
+        soft_cap=soft_cap,
+    )
+
+    output_fp16 = output.to(torch.float32) * output_scale.item()
+    output_fp16 = output_fp16.to(torch.float16)
+
+    atol, rtol = 2e-1, 2e-1
+    (
+        torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
+        f"{torch.max(torch.abs(output_fp16 - ref_output))}",
+    )

tests/kernels/core/test_fused_quant_layernorm.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
+import itertools
+
 import pytest
 import torch
 
@@ -21,7 +23,7 @@ QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
 VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
 # Avoid combinatorial explosion with full Cartesian product
 NUM_TOKENS_HIDDEN_SIZES = [
-    *[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
+    *[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],
     *[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
     *[(4096, i) for i in [1, 64, 5137]],
 ]
@@ -29,6 +31,7 @@ NUM_TOKENS_HIDDEN_SIZES = [
 ADD_RESIDUAL = [False, True]
 SCALE_UBS = [True, False]
 GROUP_SIZES = [None, [1, 64], [1, 128]]
+TMA_ALIGNMENTS = [0, 4]
 SEEDS = [0]
 CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
 
@@ -110,12 +113,21 @@ def ops_dynamic_per_token_or_block_quant(
     residual: torch.Tensor | None,
     scale_ub: torch.Tensor | None,
     group_size: list[int] | None,
+    tma_alignment: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
     if residual is not None:
         residual = residual.clone()
     if group_size is not None:
         out, scales = ops.rms_norm_per_block_quant(
-            x, weight, EPS, quant_dtype, group_size, scale_ub, residual, True
+            x,
+            weight,
+            EPS,
+            quant_dtype,
+            group_size,
+            scale_ub,
+            residual,
+            True,
+            tma_alignment,
         )
         scales = scales.contiguous()
     else:
@@ -132,9 +144,10 @@ def ops_impl(
     residual: torch.Tensor | None,
     scale_ub: torch.Tensor | None,
     group_size: list[int] | None,
+    tma_alignment: int,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
     return ops_dynamic_per_token_or_block_quant(
-        weight, x, quant_dtype, residual, scale_ub, group_size
+        weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
     )
 
 
@@ -143,7 +156,10 @@ def ops_impl(
 @pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
-@pytest.mark.parametrize("group_size", GROUP_SIZES)
+@pytest.mark.parametrize(
+    "group_size, tma_alignment",
+    [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
+)
 @pytest.mark.parametrize("seed", SEEDS)
 @pytest.mark.parametrize("device", CUDA_DEVICES)
 @torch.inference_mode()
@@ -156,6 +172,7 @@ def test_rms_norm(
     dtype: torch.dtype,
     quant_dtype: torch.dtype,
     group_size: list[int] | None,
+    tma_alignment: int,
     seed: int,
     device: str,
 ) -> None:
@@ -173,6 +190,20 @@ def test_rms_norm(
         # blockwise baseline doesn't support scale_ub
         return
 
+    if (
+        group_size is None or quant_dtype != current_platform.fp8_dtype()
+    ) and tma_alignment != 0:
+        # TMA alignment is only supported for groupwise fp8 kernels
+        return
+
+    if (
+        group_size is not None
+        and tma_alignment != 0
+        and hidden_size // group_size[1] % tma_alignment == 0
+    ):
+        # Skip tests where TMA alignment doesn't create extra padding to save time
+        return
+
     if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
         # skip
         return
@@ -196,7 +227,7 @@ def test_rms_norm(
         layer, x, quant_dtype, residual, scale_ub, group_size
     )
     ops_out, ops_scales, ops_residual = ops_impl(
-        layer.weight, x, quant_dtype, residual, scale_ub, group_size
+        layer.weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
     )
 
     assert ref_out.dtype == quant_dtype

tests/kernels/moe/test_unquantized_backend_selection.py

@@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
     UnquantizedMoeBackend,
     select_unquantized_moe_backend,
 )
+from vllm.platforms import current_platform
 
 
 @pytest.mark.parametrize(
@@ -65,6 +66,9 @@ def test_select_default_backend_by_platform(
     "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
     return_value=(True, None),
 )
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
+)
 def test_select_cuda_flashinfer_trtllm_backend(
     mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
 ):
@@ -101,6 +105,9 @@ def test_select_cuda_flashinfer_trtllm_backend(
     "vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
     return_value=(False, None),
 )
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
+)
 def test_select_cuda_flashinfer_cutlass_backend(
     mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
 ):

tests/models/multimodal/generation/test_common.py

@@ -169,13 +169,6 @@ VLM_TEST_SETTINGS = {
         auto_cls=AutoModelForImageTextToText,
         vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
         patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner,
-        vllm_runner_kwargs={
-            "attention_config": {
-                "backend": "ROCM_AITER_FA",
-            },
-        }
-        if current_platform.is_rocm()
-        else None,
         image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
         marks=[
             pytest.mark.core_model,

tests/multimodal/media/test_connector.py

@@ -7,8 +7,10 @@ import mimetypes
 import os
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
+import aiohttp
 import numpy as np
 import pytest
+import requests
 import torch
 from PIL import Image, ImageChops
 
@@ -318,3 +320,58 @@ async def test_allowed_media_domains(video_url: str, num_frames: int):
 
     with pytest.raises(ValueError):
         _, _ = await connector.fetch_video_async(disallowed_url)
+
+
+@pytest.mark.asyncio
+async def test_ssrf_bypass_backslash_in_url(local_asset_server):
+    """Verify that backslash-@ URL parsing confusion cannot bypass the
+    allowed_media_domains check (GHSA-v359-jj2v-j536).
+
+    urllib3.parse_url() and aiohttp/yarl disagree on how to parse a
+    backslash before ``@``.  urllib3 treats ``\\`` as part of the path
+    (encoding it as ``%5C``), while yarl treats it as a userinfo
+    separator, changing the effective host.  The fix normalises the URL
+    through urllib3 *before* handing it to aiohttp so both layers agree.
+    """
+    port = local_asset_server.port
+    asset = TEST_IMAGE_ASSETS[0]
+
+    # Craft the bypass payload: urllib3 sees host=127.0.0.1, but an
+    # un-patched aiohttp would see host=example.com.
+    bypass_url = f"http://127.0.0.1:{port}\\@example.com/{asset}"
+
+    connector = MediaConnector(
+        allowed_media_domains=["127.0.0.1"],
+    )
+
+    # After the fix the request is made to 127.0.0.1 (the local asset
+    # server) using the normalised URL.  The normalised path will be
+    # /%5C@example.com/<asset> which won't match any file the server
+    # knows about, so we expect an HTTP error — but crucially NOT a
+    # successful fetch from example.com.
+    with pytest.raises(requests.exceptions.HTTPError):
+        connector.fetch_image(bypass_url)
+
+    with pytest.raises(aiohttp.ClientResponseError):
+        await connector.fetch_image_async(bypass_url)
+
+
+@pytest.mark.asyncio
+async def test_ssrf_bypass_backslash_disallowed_domain():
+    """The reverse direction: even when the *attacker-controlled* host
+    appears in the urllib3-parsed hostname position the allowlist must
+    still block it.
+    """
+    # urllib3.parse_url sees host=example.com which is NOT in the
+    # allowlist, so this must be rejected before any request is made.
+    bypass_url = "https://example.com\\@safe.example.org/image.png"
+
+    connector = MediaConnector(
+        allowed_media_domains=["safe.example.org"],
+    )
+
+    with pytest.raises(ValueError, match="allowed domains"):
+        connector.fetch_image(bypass_url)
+
+    with pytest.raises(ValueError, match="allowed domains"):
+        await connector.fetch_image_async(bypass_url)

tests/reasoning/test_qwen3_reasoning_parser.py

@@ -4,46 +4,79 @@
 import pytest
 from transformers import AutoTokenizer
 
-from tests.reasoning.utils import run_reasoning_extraction
+from tests.reasoning.utils import (
+    StreamingReasoningReconstructor,
+    run_reasoning_extraction,
+    run_reasoning_extraction_streaming,
+)
 from vllm.reasoning import ReasoningParser, ReasoningParserManager
 
 parser_name = "qwen3"
 start_token = "<think>"
 end_token = "</think>"
 
-REASONING_MODEL_NAME = "Qwen/Qwen3-0.6B"
+REASONING_MODEL_NAMES = [
+    "Qwen/Qwen3-0.6B",
+    "Qwen/Qwen3.5-397B-A17B",
+    "Qwen/Qwen3-4B-Thinking-2507",
+]
 
 
-@pytest.fixture(scope="module")
-def qwen3_tokenizer():
-    return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
+@pytest.fixture(scope="module", params=REASONING_MODEL_NAMES)
+def qwen3_tokenizer(request):
+    return AutoTokenizer.from_pretrained(request.param)
 
 
-# 带 <think></think>，非stream
+# --- <think> in prompt, only </think> in output (typical) ---
+
+WITHOUT_START_TOKEN = {
+    "output": "This is a reasoning section</think>This is the rest",
+    "reasoning": "This is a reasoning section",
+    "content": "This is the rest",
+}
+WITHOUT_START_TOKEN_STREAM = {
+    "output": "This is a reasoning section</think>This is the rest",
+    "reasoning": "This is a reasoning section",
+    "content": "This is the rest",
+}
+WITHOUT_START_TOKEN_COMPLETE_REASONING = {
+    "output": "This is a reasoning section</think>",
+    "reasoning": "This is a reasoning section",
+    "content": None,
+}
+
+# --- <think> present in output (old template / edge case) ---
+
 WITH_THINK = {
     "output": "<think>This is a reasoning section</think>This is the rest",
     "reasoning": "This is a reasoning section",
     "content": "This is the rest",
 }
-# 带 <think></think>，stream
 WITH_THINK_STREAM = {
     "output": "<think>This is a reasoning section</think>This is the rest",
     "reasoning": "This is a reasoning section",
     "content": "This is the rest",
 }
-# 不带 <think></think>，非stream
+
+# --- No think tokens at all (thinking disabled) ---
+
 WITHOUT_THINK = {
     "output": "This is the rest",
     "reasoning": None,
     "content": "This is the rest",
 }
-# 不带 <think></think>，stream
+# In streaming, the parser cannot distinguish "thinking disabled" from
+# "reasoning in progress" when no think tokens have appeared yet.
+# It assumes reasoning. The serving layer handles the "thinking disabled"
+# case by checking prompt_is_reasoning_end_arr before calling the parser.
 WITHOUT_THINK_STREAM = {
     "output": "This is the rest",
-    "reasoning": None,
-    "content": "This is the rest",
+    "reasoning": "This is the rest",
+    "content": None,
 }
 
+# --- Edge cases ---
+
 COMPLETE_REASONING = {
     "output": "<think>This is a reasoning section</think>",
     "reasoning": "This is a reasoning section",
@@ -57,7 +90,7 @@ MULTILINE_REASONING = {
 ONLY_OPEN_TAG = {
     "output": "<think>This is a reasoning section",
     "reasoning": None,
-    "content": "<think>This is a reasoning section",
+    "content": "This is a reasoning section",
 }
 
 ONLY_OPEN_TAG_STREAM = {
@@ -67,6 +100,26 @@ ONLY_OPEN_TAG_STREAM = {
 }
 
 TEST_CASES = [
+    pytest.param(
+        False,
+        WITHOUT_START_TOKEN,
+        id="without_start_token",
+    ),
+    pytest.param(
+        True,
+        WITHOUT_START_TOKEN_STREAM,
+        id="without_start_token_stream",
+    ),
+    pytest.param(
+        False,
+        WITHOUT_START_TOKEN_COMPLETE_REASONING,
+        id="without_start_token_complete_reasoning",
+    ),
+    pytest.param(
+        True,
+        WITHOUT_START_TOKEN_COMPLETE_REASONING,
+        id="without_start_token_complete_reasoning_stream",
+    ),
     pytest.param(
         False,
         WITH_THINK,
@@ -140,3 +193,59 @@ def test_reasoning(
 
     assert reasoning == param_dict["reasoning"]
     assert content == param_dict["content"]
+
+
+# Multi-token delta tests: simulate real-world streaming where a single
+# delta can contain multiple tokens (e.g., speculative decoding).
+MULTI_TOKEN_DELTA_CASES = [
+    pytest.param(
+        # <think> grouped with following text in one delta
+        ["<think>This is a reasoning section", "</think>", "This is the rest"],
+        "This is a reasoning section",
+        "This is the rest",
+        id="start_token_grouped_with_text",
+    ),
+    pytest.param(
+        # </think> grouped with following content in one delta
+        ["reasoning section", "</think>This is the rest"],
+        "reasoning section",
+        "This is the rest",
+        id="end_token_grouped_with_content",
+    ),
+    pytest.param(
+        # <think> and </think> in the same delta, no content after
+        ["<think>reasoning</think>"],
+        "reasoning",
+        None,
+        id="start_and_end_in_one_delta_no_content",
+    ),
+    pytest.param(
+        # No start token, end grouped with content (Qwen3.5 style)
+        ["reasoning section", "</think>content"],
+        "reasoning section",
+        "content",
+        id="no_start_end_grouped_with_content",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "deltas, expected_reasoning, expected_content", MULTI_TOKEN_DELTA_CASES
+)
+def test_reasoning_streaming_multi_token_deltas(
+    deltas: list[str],
+    expected_reasoning: str | None,
+    expected_content: str | None,
+    qwen3_tokenizer,
+):
+    """Test that multi-token deltas don't leak <think> into reasoning."""
+    parser: ReasoningParser = ReasoningParserManager.get_reasoning_parser(parser_name)(
+        qwen3_tokenizer
+    )
+
+    reconstructor: StreamingReasoningReconstructor = run_reasoning_extraction_streaming(
+        parser, deltas
+    )
+
+    assert reconstructor.reasoning == expected_reasoning
+    assert (reconstructor.other_content or None) == expected_content

tests/renderers/test_process_multi_modal_uuids.py

@@ -6,6 +6,7 @@ import pytest
 from vllm.assets.image import ImageAsset
 from vllm.assets.video import VideoAsset
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
+from vllm.multimodal.parse import parse_mm_uuids
 from vllm.renderers.hf import HfRenderer
 from vllm.tokenizers.registry import tokenizer_args_from_config
 
@@ -45,10 +46,11 @@ def test_multi_modal_uuids_length_mismatch_raises():
     mm_uuids = {"image": ["hash_cherry"]}
 
     mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
 
     with pytest.raises(ValueError, match="must have same length as"):
-        renderer._process_mm_uuids(mm_data, mm_items, mm_uuids, "req-1")
+        renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-1")
 
 
 def test_multi_modal_uuids_missing_modality_raises():
@@ -63,10 +65,11 @@ def test_multi_modal_uuids_missing_modality_raises():
     mm_uuids = {"image": ["hash_cherry"]}
 
     mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
 
     with pytest.raises(ValueError, match="is empty but .* is missing"):
-        renderer._process_mm_uuids(mm_data, mm_items, mm_uuids, "req-2")
+        renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-2")
 
 
 @pytest.mark.parametrize(
@@ -78,7 +81,7 @@ def test_multi_modal_uuids_missing_modality_raises():
     ],
 )
 def test_multi_modal_uuids_accepts_none_and_passes_through(
-    monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
+    mm_cache_gb: float, enable_prefix_caching: bool
 ):
     renderer = _build_renderer(
         mm_cache_gb=mm_cache_gb,
@@ -94,9 +97,11 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
     mm_uuids = {"image": [None, "hash_stop"], "video": None}
 
     mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
+
     processed_mm_uuids = renderer._process_mm_uuids(
-        mm_data, mm_items, mm_uuids, "req-3"
+        mm_data, mm_data_items, mm_uuid_items, "req-3"
     )
 
     assert processed_mm_uuids == mm_uuids
@@ -111,7 +116,7 @@ def test_multi_modal_uuids_accepts_none_and_passes_through(
     ],
 )
 def test_multi_modal_uuids_accepts_empty(
-    monkeypatch, mm_cache_gb: float, enable_prefix_caching: bool
+    mm_cache_gb: float, enable_prefix_caching: bool
 ):
     renderer = _build_renderer(
         mm_cache_gb=mm_cache_gb,
@@ -124,15 +129,17 @@ def test_multi_modal_uuids_accepts_empty(
     mm_uuids = {"image": [], "video": None}  # type: ignore[var-annotated]
 
     mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
+
     processed_mm_uuids = renderer._process_mm_uuids(
-        mm_data, mm_items, mm_uuids, "req-4"
+        mm_data, mm_data_items, mm_uuid_items, "req-4"
     )
 
     assert processed_mm_uuids == mm_uuids
 
 
-def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
+def test_multi_modal_uuids_ignored_when_caching_disabled():
     # When both processor cache is 0 and prefix caching disabled, the
     # processor builds overrides from request id instead of using user UUIDs.
     renderer = _build_renderer(mm_cache_gb=0.0, enable_prefix_caching=False)
@@ -145,9 +152,11 @@ def test_multi_modal_uuids_ignored_when_caching_disabled(monkeypatch):
     mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}
 
     mm_processor = renderer.get_mm_processor()
-    mm_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+    mm_uuid_items = parse_mm_uuids(mm_uuids)
+
     processed_mm_uuids = renderer._process_mm_uuids(
-        mm_data, mm_items, mm_uuids, request_id
+        mm_data, mm_data_items, mm_uuid_items, request_id
     )
 
     # Expect request-id-based overrides are passed through

vllm/_custom_ops.py

@@ -450,15 +450,30 @@ def rms_norm_per_block_quant(
     scale_ub: torch.Tensor | None = None,
     residual: torch.Tensor | None = None,
     is_scale_transposed: bool = False,
+    tma_alignment: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert len(group_size) == 2
     output = torch.empty_like(input, dtype=quant_dtype)
     if is_scale_transposed:
-        scales = torch.empty(
-            (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
-            device=input.device,
-            dtype=torch.float32,
-        ).transpose(0, 1)
+        if tma_alignment == 0:
+            scales = torch.empty(
+                (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
+                device=input.device,
+                dtype=torch.float32,
+            ).transpose(0, 1)
+        else:
+            m = input.shape[-2]
+            sf_k = input.shape[-1] // group_size[1]
+            tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment
+            shape = input.shape[:-2] + (m, sf_k)
+            stride = (
+                (1, tma_aligned_m)
+                if input.dim() == 2
+                else (tma_aligned_m * sf_k, 1, tma_aligned_m)
+            )
+            scales = torch.empty_strided(
+                shape, stride, device=input.device, dtype=torch.float32
+            )
     else:
         scales = torch.empty(
             (input.numel() // input.shape[-1], input.shape[-1] // group_size[1]),
@@ -466,6 +481,10 @@ def rms_norm_per_block_quant(
             dtype=torch.float32,
         )
 
+    assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str(
+        tma_alignment
+    )
+
     torch.ops._C.rms_norm_per_block_quant(
         output,
         input,
@@ -2770,6 +2789,24 @@ def sm100_cutlass_mla_get_workspace_size(
     )
 
 
+def dsv3_fused_a_gemm(
+    output: torch.Tensor,
+    mat_a: torch.Tensor,
+    mat_b: torch.Tensor,
+) -> None:
+    """DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
+
+    Computes output = mat_a @ mat_b.T where:
+      mat_a: [num_tokens, 7168] row-major bf16 (hidden states)
+      mat_b: [7168, 2112] column-major bf16 (weight transposed)
+      output: [num_tokens, 2112] row-major bf16
+
+    Optimized for the DeepSeek V2/V3 QKV A-projection at small batch sizes.
+    Requires SM 9.0+ (Hopper).
+    """
+    torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b)
+
+
 if hasattr(torch.ops._C, "weight_packed_linear"):
 
     @register_fake("_C::weight_packed_linear")

vllm/compilation/passes/fusion/matcher_utils.py

@@ -292,6 +292,7 @@ class MatcherQuantFP8(MatcherCustomOp):
         has_col_major_scales: bool = False,
         is_e8m0: bool = False,
         match_rocm_aiter: bool = False,
+        is_tma_aligned: bool = False,
     ) -> None:
         if enabled is None:
             enabled = QuantFP8.enabled()
@@ -301,6 +302,7 @@ class MatcherQuantFP8(MatcherCustomOp):
         self.has_col_major_scales = has_col_major_scales
         self.is_e8m0 = is_e8m0
         self.match_rocm_aiter = match_rocm_aiter
+        self.is_tma_aligned = is_tma_aligned
 
         if match_rocm_aiter:
             assert not quant_key.scale.group_shape.is_per_tensor(), (
@@ -336,6 +338,7 @@ class MatcherQuantFP8(MatcherCustomOp):
             quant_key.scale.group_shape,
             column_major_scales=has_col_major_scales,
             use_ue8m0=is_e8m0,
+            tma_aligned_scales=self.is_tma_aligned,
             compile_native=False,
         )
 
@@ -367,8 +370,11 @@ class MatcherQuantFP8(MatcherCustomOp):
         )
 
         if self.quant_key.scale.group_shape.is_per_group():
-            assert scale is None
-            scale = self.make_scale(input, transposed=self.has_col_major_scales)
+            # for tma_aligned, the scale must be passed to forward_custom
+            # tma_aligned fusion then matches by custom op arguments
+            if not self.is_tma_aligned:
+                assert scale is None
+                scale = self.make_scale(input, transposed=self.has_col_major_scales)
 
             finfo = torch.finfo(self.quant_key.dtype)
             fp8_min = finfo.min
@@ -384,6 +390,8 @@ class MatcherQuantFP8(MatcherCustomOp):
                 fp8_min=fp8_min,
                 fp8_max=fp8_max,
                 scale_ue8m0=self.is_e8m0,
+                dummy_is_scale_transposed=self.has_col_major_scales,
+                dummy_is_tma_aligned=self.is_tma_aligned,
             )
             return result, scale
 

vllm/compilation/passes/fusion/rms_quant_fusion.py

@@ -121,6 +121,7 @@ class RMSNormQuantPattern:
         key: FusedRMSQuantKey,
         has_col_major_scales: bool = False,
         is_e8m0: bool = False,
+        is_tma_aligned: bool = False,
     ) -> None:
         self.epsilon = epsilon
         self.quant_dtype = key.quant.dtype
@@ -136,7 +137,10 @@ class RMSNormQuantPattern:
             else MatcherFusedAddRMSNorm(epsilon)
         )
         self.quant_matcher = MatcherQuantFP8(
-            key.quant, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
+            key.quant,
+            has_col_major_scales=has_col_major_scales,
+            is_e8m0=is_e8m0,
+            is_tma_aligned=is_tma_aligned,
         )
 
 
@@ -262,8 +266,9 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
         quant_dtype: torch.dtype,
         group_shape: GroupShape,
         symmetric: bool = True,
-        has_col_major_scales: bool = False,
         is_e8m0: bool = False,
+        has_col_major_scales: bool = True,
+        is_tma_aligned: bool = True,
     ) -> None:
         scale = ScaleDesc(torch.float32, False, group_shape)
         key = FusedRMSQuantKey(
@@ -271,29 +276,63 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
             quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
         )
         self.group_shape = group_shape
-        self.has_col_major_scales = has_col_major_scales
         self.is_e8m0 = is_e8m0
+        self.has_col_major_scales = has_col_major_scales
+        self.is_tma_aligned = is_tma_aligned
         super().__init__(
-            epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
+            epsilon,
+            key,
+            has_col_major_scales=has_col_major_scales,
+            is_e8m0=is_e8m0,
+            is_tma_aligned=is_tma_aligned,
         )
 
     def register(self, pm_pass: PatternMatcherPass) -> None:
         def pattern(
-            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            scale: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
-            result, scale = self.quant_matcher(result_rms)
+            result = torch.empty(
+                result_rms.shape,
+                device=result_rms.device,
+                dtype=self.quant_matcher.quant_key.dtype,
+            )
+            assert scale is not None
+            finfo = torch.finfo(self.quant_matcher.quant_key.dtype)
+            fp8_min = finfo.min
+            fp8_max = finfo.max
+
+            _, result, scale = auto_functionalized(
+                self.quant_matcher.QUANT_OP,
+                input=result_rms,
+                output_q=result,
+                output_s=scale,
+                group_size=self.quant_matcher.quant_key.scale.group_shape[1],
+                eps=1e-10,
+                fp8_min=fp8_min,
+                fp8_max=fp8_max,
+                scale_ue8m0=self.quant_matcher.is_e8m0,
+                dummy_is_scale_transposed=self.has_col_major_scales,
+                dummy_is_tma_aligned=self.is_tma_aligned,
+            )
+
             return result, residual, scale
 
         def replacement(
-            input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            residual: torch.Tensor,
+            scale: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             # In case we're matching native rms-norm, conversions might be
             # optimized out. We convert here just to be safe.
             input = input.to(dtype=self.model_dtype)
 
             result = torch.empty_like(input, dtype=self.quant_dtype)
-            scale = self.quant_matcher.make_scale(input, self.has_col_major_scales)
+
             at = auto_functionalized(
                 self.FUSED_OP,
                 result=result,
@@ -310,10 +349,12 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
             # result, residual, scale
             return at[1], at[3], at[2]
 
+        scale = self.quant_matcher.empty_f32(1, 1)
+
         pm.register_replacement(
             pattern,
             replacement,
-            self.rmsnorm_matcher.inputs(),
+            self.rmsnorm_matcher.inputs() + [scale],
             pm.fwd_only,
             pm_pass,
         )
@@ -326,8 +367,9 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
         quant_dtype: torch.dtype,
         group_shape: GroupShape,
         symmetric: bool = True,
-        has_col_major_scales: bool = False,
         is_e8m0: bool = False,
+        has_col_major_scales: bool = True,
+        is_tma_aligned: bool = True,
     ) -> None:
         scale = ScaleDesc(torch.float32, False, group_shape)
         key = FusedRMSQuantKey(
@@ -335,29 +377,55 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
             quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
         )
         self.group_shape = group_shape
+        self.has_col_major_scales = has_col_major_scales
+        self.is_tma_aligned = is_tma_aligned
         super().__init__(
-            epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0
+            epsilon,
+            key,
+            has_col_major_scales=self.has_col_major_scales,
+            is_e8m0=is_e8m0,
+            is_tma_aligned=is_tma_aligned,
         )
 
     def register(self, pm_pass: PatternMatcherPass) -> None:
         def pattern(
-            input: torch.Tensor, weight: torch.Tensor
+            input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor]:
             result_rms = self.rmsnorm_matcher(input, weight)
-            result, scale = self.quant_matcher(result_rms)
+            result = torch.empty(
+                result_rms.shape,
+                device=result_rms.device,
+                dtype=self.quant_matcher.quant_key.dtype,
+            )
+            assert scale is not None
+            finfo = torch.finfo(self.quant_matcher.quant_key.dtype)
+            fp8_min = finfo.min
+            fp8_max = finfo.max
+
+            _, result, scale = auto_functionalized(
+                self.quant_matcher.QUANT_OP,
+                input=result_rms,
+                output_q=result,
+                output_s=scale,
+                group_size=self.quant_matcher.quant_key.scale.group_shape[1],
+                eps=1e-10,
+                fp8_min=fp8_min,
+                fp8_max=fp8_max,
+                scale_ue8m0=self.quant_matcher.is_e8m0,
+                dummy_is_scale_transposed=self.has_col_major_scales,
+                dummy_is_tma_aligned=self.is_tma_aligned,
+            )
+
             return result, scale
 
         def replacement(
-            input: torch.Tensor, weight: torch.Tensor
+            input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
         ) -> tuple[torch.Tensor, torch.Tensor]:
             # In case we're matching native rms-norm, conversions might be
             # optimized out. We convert here just to be safe.
             input = input.to(dtype=self.model_dtype)
 
             result = torch.empty_like(input, dtype=self.quant_dtype)
-            scale = self.quant_matcher.make_scale(
-                input, transposed=self.quant_matcher.has_col_major_scales
-            )
             at = auto_functionalized(
                 self.FUSED_OP,
                 result=result,
@@ -368,16 +436,18 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
                 scale_ub=None,
                 residual=None,
                 group_size=self.group_shape[1],
-                is_scale_transposed=self.quant_matcher.has_col_major_scales,
+                is_scale_transposed=self.has_col_major_scales,
             )
 
             # result, scale
             return at[1], at[2]
 
+        scale = self.quant_matcher.empty_f32(1, 1)
+
         pm.register_replacement(
             pattern,
             replacement,
-            self.rmsnorm_matcher.inputs(),
+            self.rmsnorm_matcher.inputs() + [scale],
             pm.fwd_only,
             pm_pass,
         )
@@ -532,23 +602,26 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
                 for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:
                     for has_col_major_scales in [True, False]:
                         for is_e8m0 in [True, False]:
-                            # Fuse fused_add_rms_norm + fp8 group quant
-                            FusedAddRMSNormGroupQuantPattern(
-                                epsilon,
-                                FP8_DTYPE,
-                                group_shape=group_shape,
-                                has_col_major_scales=has_col_major_scales,
-                                is_e8m0=is_e8m0,
-                            ).register(self.patterns)
+                            for is_tma_aligned in [False, True]:
+                                # Fuse fused_add_rms_norm + fp8 group quant
+                                FusedAddRMSNormGroupQuantPattern(
+                                    epsilon,
+                                    FP8_DTYPE,
+                                    group_shape=group_shape,
+                                    is_e8m0=is_e8m0,
+                                    has_col_major_scales=has_col_major_scales,
+                                    is_tma_aligned=is_tma_aligned,
+                                ).register(self.patterns)
 
-                            # Fuse rms_norm + fp8 group quant
-                            RMSNormGroupQuantPattern(
-                                epsilon,
-                                FP8_DTYPE,
-                                group_shape=group_shape,
-                                has_col_major_scales=has_col_major_scales,
-                                is_e8m0=is_e8m0,
-                            ).register(self.patterns)
+                                # Fuse rms_norm + fp8 group quant
+                                RMSNormGroupQuantPattern(
+                                    epsilon,
+                                    FP8_DTYPE,
+                                    group_shape=group_shape,
+                                    is_e8m0=is_e8m0,
+                                    has_col_major_scales=has_col_major_scales,
+                                    is_tma_aligned=is_tma_aligned,
+                                ).register(self.patterns)
 
         self.dump_patterns(config, self.patterns)
 

vllm/config/vllm.py

@@ -95,11 +95,16 @@ def enable_norm_fusion(cfg: "VllmConfig") -> bool:
 
 
 def enable_act_fusion(cfg: "VllmConfig") -> bool:
-    """Enable if either SiLU+Mul or quant FP8 custom op is active;
-    otherwise Inductor handles fusion."""
-    return cfg.compilation_config.is_custom_op_enabled(
-        "silu_and_mul"
-    ) or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
+    """
+    Enable if either SiLU+Mul or quant FP8 custom op is active;
+    otherwise Inductor handles fusion.
+    Also enable for FP4 models as FP4 quant is always custom so Inductor cannot fuse it.
+    """
+    return (
+        cfg.compilation_config.is_custom_op_enabled("silu_and_mul")
+        or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
+        or (cfg.model_config is not None and cfg.model_config.is_nvfp4_quantized())
+    )
 
 
 def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:

vllm/entrypoints/llm.py

@@ -519,7 +519,6 @@ class LLM:
             ),
             params=seq_params,
             lora_requests=seq_lora_requests,
-            tokenization_kwargs=tokenization_kwargs,
             priorities=seq_priority,
         )
 
@@ -1813,7 +1812,6 @@ class LLM:
             params=seq_params,
             use_tqdm=use_tqdm,
             lora_requests=seq_lora_requests,
-            tokenization_kwargs=tokenization_kwargs,
             priorities=seq_priority,
         )
 
@@ -1872,7 +1870,6 @@ class LLM:
             params=seq_params,
             lora_requests=seq_lora_requests,
             use_tqdm=use_tqdm,
-            tokenization_kwargs=tokenization_kwargs,
         )
 
     def _render_and_run_requests(
@@ -1881,7 +1878,6 @@ class LLM:
         params: Sequence[SamplingParams | PoolingParams],
         *,
         lora_requests: Sequence[LoRARequest | None] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
         priorities: Sequence[int] | None = None,
         use_tqdm: bool | Callable[..., tqdm] = True,
     ):
@@ -1899,7 +1895,6 @@ class LLM:
             prompts=prompts,
             params=params,
             lora_requests=lora_requests,
-            tokenization_kwargs=tokenization_kwargs,
             priorities=priorities,
         )
 
@@ -1911,7 +1906,6 @@ class LLM:
         params: Sequence[SamplingParams | PoolingParams],
         *,
         lora_requests: Sequence[LoRARequest | None] | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
         priorities: Sequence[int] | None = None,
     ) -> list[str]:
         added_request_ids: list[str] = []
@@ -1922,7 +1916,6 @@ class LLM:
                     prompt,
                     params[i],
                     lora_request=None if lora_requests is None else lora_requests[i],
-                    tokenization_kwargs=tokenization_kwargs,
                     priority=0 if priorities is None else priorities[i],
                 )
                 added_request_ids.append(request_id)
@@ -1938,7 +1931,6 @@ class LLM:
         prompt: ProcessorInputs,
         params: SamplingParams | PoolingParams,
         lora_request: LoRARequest | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
         priority: int = 0,
     ) -> str:
         if isinstance(params, SamplingParams):
@@ -1947,27 +1939,11 @@ class LLM:
 
         request_id = str(next(self.request_counter))
 
-        if params.truncate_prompt_tokens is not None:
-            params_type = type(params).__name__
-            warnings.warn(
-                f"The `truncate_prompt_tokens` parameter in `{params_type}` "
-                "is deprecated and will be removed in v0.16. "
-                "Please pass it via `tokenization_kwargs` instead.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-
-            tokenization_kwargs = merge_kwargs(
-                tokenization_kwargs,
-                dict(truncate_prompt_tokens=params.truncate_prompt_tokens),
-            )
-
         return self.llm_engine.add_request(
             request_id,
             prompt,
             params,
             lora_request=lora_request,
-            tokenization_kwargs=tokenization_kwargs,
             priority=priority,
         )
 

vllm/entrypoints/openai/chat_completion/serving.py

@@ -900,6 +900,17 @@ class OpenAIServingChat(OpenAIServing):
                         harmony_tools_streamed[i] |= tools_streamed_flag
                     # handle streaming deltas for tools with named tool_choice
                     elif tool_choice_function_name:
+                        # When encountering think end id in prompt_token_ids
+                        # i.e {"enable_thinking": False},
+                        # check BEFORE calling the parser to avoid a spurious
+                        # reasoning delta on the first chunk.
+                        if (
+                            reasoning_parser
+                            and not reasoning_end_arr[i]
+                            and prompt_is_reasoning_end_arr[i]
+                        ):
+                            reasoning_end_arr[i] = True
+
                         if (
                             reasoning_parser
                             and not reasoning_end_arr[i]
@@ -918,16 +929,11 @@ class OpenAIServingChat(OpenAIServing):
                                     output.token_ids,
                                 )
                             )
-                            # When encountering think end id in delta_token_ids
-                            # or think end id in prompt_token_ids
-                            # i.e {"enable_thinking": False},
+                            # When encountering think end id in delta_token_ids,
                             # set reasoning status to end.
                             # Only keep 'content', remove 'reasoning'.
-                            if (
-                                reasoning_parser.is_reasoning_end(
-                                    as_list(output.token_ids)
-                                )
-                                or prompt_is_reasoning_end_arr[i]
+                            if reasoning_parser.is_reasoning_end(
+                                as_list(output.token_ids)
                             ):
                                 reasoning_end_arr[i] = True
                                 if delta_message and delta_message.content:
@@ -1116,14 +1122,23 @@ class OpenAIServingChat(OpenAIServing):
 
                     # when only reasoning
                     elif reasoning_parser:
-                        delta_message = reasoning_parser.extract_reasoning_streaming(
-                            previous_text,
-                            current_text,
-                            delta_text,
-                            previous_token_ids,
-                            current_token_ids,
-                            output.token_ids,
-                        )
+                        # When encountering think end id in prompt_token_ids
+                        # i.e {"enable_thinking": False},
+                        # set reasoning status to end.
+                        # Route all generated tokens as content directly.
+                        if prompt_is_reasoning_end_arr[i]:
+                            delta_message = DeltaMessage(content=delta_text)
+                        else:
+                            delta_message = (
+                                reasoning_parser.extract_reasoning_streaming(
+                                    previous_text,
+                                    current_text,
+                                    delta_text,
+                                    previous_token_ids,
+                                    current_token_ids,
+                                    output.token_ids,
+                                )
+                            )
                     # handle streaming just a content delta
                     else:
                         delta_message = DeltaMessage(content=delta_text)

vllm/inputs/preprocess.py

@@ -91,7 +91,7 @@ class InputPreprocessor:
         self,
         prompt: str | list[int],
         mm_data: MultiModalDataDict,
-        mm_processor_kwargs: Mapping[str, object] | None,
+        mm_processor_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: dict[str, Any] | None = None,
         *,
         mm_uuids: MultiModalUUIDDict | None = None,
@@ -103,9 +103,9 @@ class InputPreprocessor:
         return self.renderer._process_multimodal(
             prompt,
             mm_data,
+            mm_uuids=mm_uuids,
             mm_processor_kwargs=mm_processor_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
     def _process_embeds(
@@ -144,7 +144,7 @@ class InputPreprocessor:
             inputs = self._process_multimodal(
                 prompt_token_ids,
                 multi_modal_data,
-                parsed_content.get("mm_processor_kwargs") or {},
+                parsed_content.get("mm_processor_kwargs"),
                 tokenization_kwargs=tokenization_kwargs,
                 mm_uuids=parsed_content.get("multi_modal_uuids"),
             )

vllm/model_executor/layers/fused_moe/cpu_fused_moe.py

@@ -402,7 +402,7 @@ class CPUFusedMOE:
             input,
             topk_weights,
             topk_ids,
-            activation,
+            activation.value,
             global_num_experts,
             skip_weighted,
         )

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

@@ -381,7 +381,6 @@ class AiterExperts(mk.FusedMoEPermuteExpertsUnpermute):
         # TODO(rob): rocm_aiter_fused_experts uses self.quant_config's
         # a_scales for static quantization. Update this to fit better
         # with the interface once all quant integrations are complete.
-        assert a2_scale == self.quant_config.a2_scale
 
         if expert_tokens_meta is not None:
             num_local_tokens = expert_tokens_meta.expert_num_tokens

vllm/model_executor/layers/linear.py

@@ -685,8 +685,13 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         self,
         param: Parameter,
         loaded_weight: torch.Tensor,
-        loaded_shard_id: int | None = None,
+        loaded_shard_id: tuple[int, ...] | int | None = None,
     ):
+        if isinstance(loaded_shard_id, tuple):
+            raise NotImplementedError(
+                "Shard id with multiple indices is not supported in weight_loader, "
+                "please use weight_loader_v2 instead."
+            )
         # Special case for GGUF
         # initialize GGUF param after we know the quantize type
         is_gguf_weight = getattr(param, "is_gguf_weight", False)
@@ -770,6 +775,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         if output_dim is not None:
             shard_offset = sum(self.output_sizes[:loaded_shard_id])
             shard_size = self.output_sizes[loaded_shard_id]
+            shard_offset //= self.tp_size
+            shard_size //= self.tp_size
 
             if isinstance(param, BlockQuantScaleParameter):
                 weight_block_size = getattr(self, "weight_block_size", None)
@@ -777,9 +784,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                     weight_block_size, shard_size, shard_offset
                 )
 
-            shard_offset //= self.tp_size
-            shard_size //= self.tp_size
-
             # Special case for quantization.
             # If quantized, we need to adjust the offset and size to account
             # for the packing.
@@ -825,7 +829,10 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         param_data.copy_(loaded_weight)
 
     def _load_fused_module_from_checkpoint(
-        self, param: BasevLLMParameter, loaded_weight: torch.Tensor
+        self,
+        param: BasevLLMParameter,
+        loaded_weight: torch.Tensor,
+        output_sizes: list[int] | None = None,
     ):
         """
         Handle special case for models where MLP layers are already
@@ -839,7 +846,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
 
         current_shard_offset = 0
         shard_offsets: list[tuple[int, int, int]] = []
-        for i, output_size in enumerate(self.output_sizes):
+        output_sizes = output_sizes or self.output_sizes
+        for i, output_size in enumerate(output_sizes):
             shard_offsets.append((i, current_shard_offset, output_size))
             current_shard_offset += output_size
 
@@ -864,23 +872,38 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
         self,
         param: BasevLLMParameter,
         loaded_weight: torch.Tensor,
-        loaded_shard_id: int | None = None,
+        loaded_shard_id: tuple[int, ...] | int | None = None,
     ):
-        if loaded_shard_id is None:
+        if loaded_shard_id is None or isinstance(loaded_shard_id, tuple):
             if isinstance(param, PerTensorScaleParameter):
                 param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0)
                 return
             elif type(param) in (RowvLLMParameter, BasevLLMParameter):
                 param.load_merged_column_weight(loaded_weight=loaded_weight)
                 return
+            output_sizes = (
+                [self.output_sizes[idx] for idx in loaded_shard_id]
+                if loaded_shard_id
+                else None
+            )
+            if isinstance(param, BlockQuantScaleParameter):
+                weight_block_size = getattr(self, "weight_block_size", None)
+                output_sizes = [
+                    adjust_block_scale_shard(weight_block_size, size, 0)[0]
+                    for size in (output_sizes or self.output_sizes)
+                ]
             # TODO: @dsikka - move to parameter.py
-            self._load_fused_module_from_checkpoint(param, loaded_weight)
+            self._load_fused_module_from_checkpoint(
+                param, loaded_weight, output_sizes=output_sizes
+            )
             return
 
         assert loaded_shard_id < len(self.output_sizes)
 
         shard_offset = sum(self.output_sizes[:loaded_shard_id])
         shard_size = self.output_sizes[loaded_shard_id]
+        shard_offset //= self.tp_size
+        shard_size //= self.tp_size
 
         if isinstance(param, BlockQuantScaleParameter):
             weight_block_size = getattr(self, "weight_block_size", None)
@@ -888,9 +911,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                 weight_block_size, shard_size, shard_offset
             )
 
-        shard_offset //= self.tp_size
-        shard_size //= self.tp_size
-
         param.load_merged_column_weight(
             loaded_weight=loaded_weight,
             shard_id=loaded_shard_id,

vllm/model_executor/layers/mla.py

@@ -129,6 +129,7 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
             assert self.q_b_proj is not None, (
                 "q_b_proj is required when q_lora_rank is not None"
             )
+
             qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]
             q_c, kv_lora = qkv_lora.split(
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

vllm/model_executor/layers/quantization/base_config.py

@@ -18,6 +18,11 @@ else:
 class QuantizeMethodBase(ABC):
     """Base class for different quantized methods."""
 
+    # Whether this method creates weights on meta device for online quantization.
+    # When True, weights are created on meta device and quantized layer-wise
+    # in process_weights_after_loading, reducing peak memory during loading.
+    uses_meta_device: bool = False
+
     @abstractmethod
     def create_weights(
         self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs

vllm/model_executor/layers/quantization/fp8.py

@@ -527,6 +527,8 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
     """Online version of Fp8LinearMethod, loads the fp16/bf16 checkpoint
     and quantized the weights during loading."""
 
+    uses_meta_device: bool = True
+
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -1039,6 +1041,8 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
         quant_config: The quantization config.
     """
 
+    uses_meta_device: bool = True
+
     def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
         super().__init__(quant_config, layer)
         assert not quant_config.is_checkpoint_fp8_serialized

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

@@ -122,6 +122,8 @@ def is_supported_config_trtllm(
         return False, _make_reason("routing method")
     elif activation_format != mk.FusedMoEActivationFormat.Standard:
         return False, _make_reason("activation format")
+    elif moe_config.hidden_dim % 512 != 0:
+        return False, _make_reason("hidden_dim must be divisible by 512")
 
     return True, None
 

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

@@ -455,4 +455,15 @@ def prepare_fp8_moe_layer_for_fi(
             w2_input_scale=w2_input_scale,
         )
 
+    # Clamp block scales to avoid NaN from the FlashInfer CUTLASS kernel.
+    # Some FP8 models have near-zero block scales (~1e-23) for dead/unused
+    # experts. The CUTLASS kernel doesn't handle these correctly on Hopper
+    # (SM 9.0), producing NaN instead of near-zero output. Clamping to a
+    # small minimum prevents this without affecting model accuracy since
+    # these experts' effective weights are already zero.
+    if block_quant:
+        _FI_CUTLASS_MIN_BLOCK_SCALE = 1e-10
+        w13_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE)
+        w2_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE)
+
     return w13, w2, w13_scale

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -924,7 +924,16 @@ def per_token_group_quant_fp8(
     # TODO(bnell): this causes some fp8 moe test to fail.
     if current_platform.is_cuda() and x.is_contiguous():
         torch.ops._C.per_token_group_fp8_quant(
-            x, x_q, x_s, group_size, eps, fp8_min, fp8_max, use_ue8m0
+            x,
+            x_q,
+            x_s,
+            group_size,
+            eps,
+            fp8_min,
+            fp8_max,
+            use_ue8m0,
+            column_major_scales,
+            tma_aligned_scales,
         )
         return x_q, x_s
 

vllm/model_executor/model_loader/weight_utils.py

@@ -1092,16 +1092,20 @@ def initialize_dummy_weights(
     is fixed, the random values generated by this function only depends on
     the parameter's number of elements and its data type.
     """
-    # TODO(future PR): make the check below more generic as more online
-    # quant backends are added
-    is_fp8_py_quant = model_config.quantization == "fp8"
+
+    # Check if any module uses online quantization with meta device weights.
+    # If so, we'll skip initializing params on meta device since they'll be
+    # handled in `process_weights_after_loading`.
+    def uses_meta_device(module: torch.nn.Module) -> bool:
+        quant_method = getattr(module, "quant_method", None)
+        return getattr(quant_method, "uses_meta_device", False)
+
+    has_online_quant = any(uses_meta_device(m) for m in model.modules())
 
     for param in model.state_dict().values():
-        if is_fp8_py_quant and param.device == torch.device("meta"):
-            # for fp8.py's online quantization, dummy weight init will happen
-            # in `process_weights_after_loading`.
-            # TODO(future PR): consider refactoring dummy model init to compose
-            # better with online quantization
+        if has_online_quant and param.device == torch.device("meta"):
+            # For online quantization, weights are created on meta device and
+            # dummy weight init will happen in `process_weights_after_loading`.
             continue
 
         initialize_single_dummy_weight(param, low, high, seed)

vllm/model_executor/models/clip.py

@@ -36,9 +36,13 @@ from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalInputs,
     MultiModalKwargsItems,
-    MultiModalUUIDDict,
 )
-from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
+from vllm.multimodal.parse import (
+    ImageProcessorItems,
+    ImageSize,
+    MultiModalDataItems,
+    MultiModalUUIDItems,
+)
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
@@ -203,10 +207,9 @@ class CLIPMultiModalProcessor(BaseMultiModalProcessor[CLIPProcessingInfo]):
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalInputs:
         if mm_items:
             if isinstance(prompt, str):
@@ -235,9 +238,9 @@ class CLIPMultiModalProcessor(BaseMultiModalProcessor[CLIPProcessingInfo]):
         return super().apply(
             prompt=prompt,
             mm_items=mm_items,
+            mm_uuid_items=mm_uuid_items,
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
     def _hf_processor_applies_updates(

vllm/model_executor/models/deepseek_v2.py

@@ -32,6 +32,7 @@ import torch
 from torch import nn
 from transformers import DeepseekV2Config, DeepseekV3Config
 
+import vllm._custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
@@ -711,6 +712,64 @@ class Indexer(nn.Module):
         return self.indexer_op(hidden_states, q_fp8, k, weights)
 
 
+class DeepSeekV2FusedQkvAProj(MergedColumnParallelLinear):
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
+        super().__init__(
+            input_size,
+            output_size,
+            bias=False,
+            quant_config=quant_config,
+            disable_tp=True,
+            prefix=f"{prefix}.kv_a_proj_with_mqa",
+        )
+
+        # Check if the DeepSeek V3 fused A GEMM kernel can be used.
+        # This kernel supports PDL and is optimized for low batch size.
+        self._use_min_latency_gemm = (
+            hasattr(self, "weight")
+            and self.weight.dtype == torch.bfloat16
+            and self.weight.shape[0] == 2112
+            and self.weight.shape[1] == 7168
+            and current_platform.is_cuda()
+            and (
+                current_platform.is_device_capability(90)
+                or current_platform.is_device_capability_family(100)
+            )
+        )
+
+    def forward(
+        self,
+        input_,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.nn.Parameter | None]:
+        num_tokens = input_.shape[0]
+        if self._use_min_latency_gemm and (0 < num_tokens <= 16):
+            output = torch.empty(
+                num_tokens,
+                2112,
+                dtype=torch.bfloat16,
+                device=input_.device,
+            )
+            ops.dsv3_fused_a_gemm(
+                output,
+                input_,
+                self.weight.T,
+            )
+            if not self.return_bias:
+                return output
+            output_bias = self.bias if self.skip_bias_add else None
+            return output, output_bias
+        else:
+            # Fallback to the standard forward method when
+            # the fused A GEMM kernel cannot be used.
+            return super().forward(input_)
+
+
 class DeepseekV2MLAAttention(nn.Module):
     """
     Main reference: DeepseekV2 paper, and FlashInfer Implementation
@@ -756,13 +815,11 @@ class DeepseekV2MLAAttention(nn.Module):
         self.max_position_embeddings = max_position_embeddings
 
         if self.q_lora_rank is not None:
-            self.fused_qkv_a_proj = MergedColumnParallelLinear(
+            self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProj(
                 self.hidden_size,
                 [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
-                bias=False,
                 quant_config=quant_config,
                 prefix=f"{prefix}.fused_qkv_a_proj",
-                disable_tp=True,
             )
         else:
             self.kv_a_proj_with_mqa = ReplicatedLinear(

vllm/model_executor/models/deepseek_vl2.py

@@ -24,13 +24,13 @@ from vllm.multimodal.inputs import (
     MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalKwargsItems,
-    MultiModalUUIDDict,
 )
 from vllm.multimodal.parse import (
     ImageEmbeddingItems,
     ImageProcessorItems,
     ImageSize,
     MultiModalDataItems,
+    MultiModalUUIDItems,
 )
 from vllm.multimodal.processing import BaseDummyInputsBuilder
 from vllm.multimodal.processing.processor import (
@@ -313,9 +313,9 @@ class DeepseekVL2MultiModalProcessor(
         self,
         prompt: str | list[int],
         mm_data_items: MultiModalDataItems,
+        mm_uuid_items: MultiModalUUIDItems | None,
         hf_processor_mm_kwargs: Mapping[str, object],
         tokenization_kwargs: Mapping[str, object],
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
         # The processor logic is different for len(images) <= 2 vs > 2
         # Since the processing cache assumes that the processor output is
@@ -325,17 +325,17 @@ class DeepseekVL2MultiModalProcessor(
             return self._apply_hf_processor(
                 prompt=prompt,
                 mm_data_items=mm_data_items,
+                mm_uuid_items=mm_uuid_items,
                 hf_processor_mm_kwargs=hf_processor_mm_kwargs,
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
 
         return super()._cached_apply_hf_processor(
             prompt=prompt,
             mm_data_items=mm_data_items,
+            mm_uuid_items=mm_uuid_items,
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
 

vllm/model_executor/models/h2ovl.py

@@ -16,11 +16,12 @@ from transformers import PretrainedConfig
 
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.multimodal.inputs import MultiModalKwargsItems, MultiModalUUIDDict
+from vllm.multimodal.inputs import MultiModalKwargsItems
 from vllm.multimodal.parse import (
     ImageEmbeddingItems,
     ImageProcessorItems,
     MultiModalDataItems,
+    MultiModalUUIDItems,
 )
 from vllm.multimodal.processing.processor import (
     MultiModalProcessingInfo,
@@ -491,9 +492,9 @@ class H2OVLMultiModalProcessor(BaseInternVLMultiModalProcessor[H2OVLProcessingIn
         self,
         prompt: str | list[int],
         mm_data_items: MultiModalDataItems,
+        mm_uuid_items: MultiModalUUIDItems | None,
         hf_processor_mm_kwargs: Mapping[str, object],
         tokenization_kwargs: Mapping[str, object],
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
         # The processor logic is different for len(images) <= 1 vs > 1
         # Since the processing cache assumes that the processor output is
@@ -503,17 +504,17 @@ class H2OVLMultiModalProcessor(BaseInternVLMultiModalProcessor[H2OVLProcessingIn
             return self._apply_hf_processor(
                 prompt=prompt,
                 mm_data_items=mm_data_items,
+                mm_uuid_items=mm_uuid_items,
                 hf_processor_mm_kwargs=hf_processor_mm_kwargs,
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
 
         return super()._cached_apply_hf_processor(
             prompt=prompt,
             mm_data_items=mm_data_items,
+            mm_uuid_items=mm_uuid_items,
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
 

vllm/model_executor/models/llava.py

@@ -30,7 +30,6 @@ from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalInputs,
     MultiModalKwargsItems,
-    MultiModalUUIDDict,
     mm_inputs,
 )
 from vllm.multimodal.parse import (
@@ -38,6 +37,7 @@ from vllm.multimodal.parse import (
     ImageProcessorItems,
     ImageSize,
     MultiModalDataItems,
+    MultiModalUUIDItems,
 )
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
@@ -773,9 +773,9 @@ class MantisMultiModalProcessor(LlavaMultiModalProcessor):
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalInputs:
         hf_config = self.info.get_hf_config()
         image_token_id = hf_config.image_token_index
@@ -789,9 +789,9 @@ class MantisMultiModalProcessor(LlavaMultiModalProcessor):
         result = super().apply(
             prompt,
             mm_items,
-            hf_processor_mm_kwargs,
-            tokenization_kwargs,
-            mm_uuids=mm_uuids,
+            mm_uuid_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
         )
 
         mm_item_counts = mm_items.get_all_counts()

vllm/model_executor/models/paligemma.py

@@ -16,12 +16,12 @@ from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalInputs,
     MultiModalKwargsItems,
-    MultiModalUUIDDict,
 )
 from vllm.multimodal.parse import (
     ImageEmbeddingItems,
     ImageProcessorItems,
     MultiModalDataItems,
+    MultiModalUUIDItems,
 )
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
@@ -231,16 +231,16 @@ class PaliGemmaMultiModalProcessor(BaseMultiModalProcessor[PaliGemmaProcessingIn
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalInputs:
         mm_inputs = super().apply(
             prompt,
             mm_items,
-            hf_processor_mm_kwargs,
-            tokenization_kwargs,
-            mm_uuids=mm_uuids,
+            mm_uuid_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
         )
         prompt_token_ids = mm_inputs["prompt_token_ids"]
 

vllm/model_executor/models/pixtral.py

@@ -44,10 +44,14 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargsItems
 from vllm.multimodal.inputs import (
     MultiModalDataDict,
     MultiModalFieldConfig,
-    MultiModalUUIDDict,
     NestedTensors,
 )
-from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
+from vllm.multimodal.parse import (
+    ImageProcessorItems,
+    ImageSize,
+    MultiModalDataItems,
+    MultiModalUUIDItems,
+)
 from vllm.multimodal.processing import BaseDummyInputsBuilder, ProcessorInputs
 from vllm.multimodal.processing.processor import (
     BaseMultiModalProcessor,
@@ -344,16 +348,16 @@ class PixtralMultiModalProcessor(BaseMultiModalProcessor[PixtralProcessingInfo])
         self,
         prompt: str | list[int],
         mm_data_items: MultiModalDataItems,
+        mm_uuid_items: MultiModalUUIDItems | None,
         hf_processor_mm_kwargs: Mapping[str, object],
         tokenization_kwargs: Mapping[str, object],
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
         prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(
             prompt=prompt,
             mm_data_items=mm_data_items,
+            mm_uuid_items=mm_uuid_items,
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
         # NOTE: The tokens are already inserted by the chat template

vllm/model_executor/models/qwen3_5.py

@@ -30,36 +30,20 @@ from collections.abc import Callable, Iterable
 import torch
 from einops import rearrange
 from torch import nn
-from transformers.activations import ACT2FN
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import (
-    CacheConfig,
-    ModelConfig,
-    SpeculativeConfig,
     VllmConfig,
-    get_current_vllm_config,
 )
 from vllm.distributed import (
-    divide,
     get_pp_group,
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
 )
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import (
     GemmaRMSNorm as Qwen3_5RMSNorm,
 )
-from vllm.model_executor.layers.layernorm import RMSNormGated
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear,
-    MergedColumnParallelLinear,
-    RowParallelLinear,
-)
+from vllm.model_executor.layers.linear import MergedColumnParallelLinear
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
-from vllm.model_executor.layers.mamba.mamba_mixer2 import (
-    mamba_v2_sharded_weight_loader,
-)
 from vllm.model_executor.layers.mamba.mamba_utils import (
     MambaStateCopyFunc,
     MambaStateCopyFuncCalculator,
@@ -73,11 +57,8 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 )
 from vllm.model_executor.model_loader.weight_utils import (
     default_weight_loader,
-    sharded_weight_loader,
 )
-from vllm.model_executor.utils import set_weight_attrs
 from vllm.multimodal import MULTIMODAL_REGISTRY
-from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.qwen3_5 import (
     Qwen3_5Config,
@@ -99,7 +80,6 @@ from .interfaces import (
 )
 from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
 from .qwen3_next import (
-    ChunkGatedDeltaRule,
     Qwen3NextAttention,
     Qwen3NextDecoderLayer,
     Qwen3NextGatedDeltaNet,
@@ -139,154 +119,31 @@ class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo):
 
 
 class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
-    def __init__(
-        self,
-        config: Qwen3_5TextConfig | Qwen3_5MoeTextConfig,
-        model_config: ModelConfig | None = None,
-        cache_config: CacheConfig | None = None,
-        quant_config: QuantizationConfig | None = None,
-        speculative_config: SpeculativeConfig | None = None,
-        prefix: str = "",
-    ) -> None:
-        super(Qwen3NextGatedDeltaNet, self).__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
-        self.tp_rank = get_tensor_model_parallel_rank()
-        self.hidden_size = config.hidden_size
-        self.num_v_heads = config.linear_num_value_heads
-        self.num_k_heads = config.linear_num_key_heads
-        self.head_k_dim = config.linear_key_head_dim
-        self.head_v_dim = config.linear_value_head_dim
-        self.key_dim = self.head_k_dim * self.num_k_heads
-        self.value_dim = self.head_v_dim * self.num_v_heads
-
-        self.conv_kernel_size = config.linear_conv_kernel_dim
-        self.layer_idx = extract_layer_index(prefix)
-        self.activation = config.hidden_act
-        self.act = ACT2FN[config.hidden_act]
-        self.layer_norm_epsilon = config.rms_norm_eps
-        self.prefix = prefix
-
-        self.config = config
-        self.model_config = model_config
-        self.cache_config = cache_config
-        self.quant_config = quant_config
-        self.speculative_config = speculative_config
-        self.num_spec = (
-            self.speculative_config.num_speculative_tokens
-            if self.speculative_config
-            else 0
-        )
-
-        # QKV
-        self.conv_dim = self.key_dim * 2 + self.value_dim
-        self.conv1d = ColumnParallelLinear(
-            input_size=self.conv_kernel_size,
-            output_size=self.conv_dim,
-            bias=False,
-            prefix=f"{prefix}.conv1d",
-        )
-        self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
-
-        self.in_proj_qkv = MergedColumnParallelLinear(
-            input_size=self.hidden_size,
-            output_sizes=[self.key_dim, self.key_dim, self.value_dim],
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.in_proj_qkv",
-        )
-        self.in_proj_z = ColumnParallelLinear(
-            input_size=self.hidden_size,
-            output_size=self.value_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.in_proj_z",
-        )
-        self.in_proj_b = ColumnParallelLinear(
-            input_size=self.hidden_size,
-            output_size=self.num_v_heads,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.in_proj_b",
-        )
-        self.in_proj_a = ColumnParallelLinear(
-            input_size=self.hidden_size,
-            output_size=self.num_v_heads,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.in_proj_a",
-        )
-
-        query_key_settings = (self.key_dim, 0, False)
-        value_settings = (self.value_dim, 0, False)
-
-        delattr(self.conv1d.weight, "weight_loader")
-        set_weight_attrs(
-            self.conv1d.weight,
-            {
-                "weight_loader": mamba_v2_sharded_weight_loader(
-                    [
-                        query_key_settings,
-                        query_key_settings,
-                        value_settings,
-                    ],
-                    self.tp_size,
-                    self.tp_rank,
-                )
-            },
-        )
-
-        # selective projection used to make dt, B and C input dependant
-
-        # time step projection (discretization)
-        # instantiate once and copy inv_dt in init_weights of PretrainedModel
-        self.dt_bias = nn.Parameter(
-            torch.ones(self.num_v_heads // self.tp_size),
-        )
-        self.A_log = nn.Parameter(
-            torch.empty(
-                divide(self.num_v_heads, self.tp_size),
-            )
-        )
-
-        set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
-        set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
-
-        self.norm = RMSNormGated(
-            self.head_v_dim,
-            eps=self.layer_norm_epsilon,
-            group_size=None,
-            norm_before_gate=True,
-            device=current_platform.current_device(),
-            dtype=config.dtype,
-        )
-
-        self.out_proj = RowParallelLinear(
-            self.value_dim,
-            self.hidden_size,
-            bias=False,
-            input_is_parallel=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.out_proj",
-        )
-
-        self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
-
-        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
-
     def fix_query_key_value_ordering(
         self,
-        mixed_qkv,
-        z,
-        b,
-        a,
+        mixed_qkvz: torch.Tensor,
+        mixed_ba: torch.Tensor,
     ):
         raise NotImplementedError(
             "Qwen3.5 Series dont need to fix query key value ordering"
         )
 
+    def create_qkvz_proj(
+        self,
+        hidden_size: int,
+        key_dim: int,
+        value_dim: int,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> MergedColumnParallelLinear:
+        return MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[key_dim, key_dim, value_dim, value_dim],
+            bias=False,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -303,11 +160,13 @@ class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
         # ============================================================
         # Part 1: Input Projection
         # ============================================================
-        mixed_qkv, _ = self.in_proj_qkv(hidden_states)
-        z, _ = self.in_proj_z(hidden_states)
+        mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
+        qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
+        z_size = self.value_dim // self.tp_size
+        mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
         z = z.reshape(z.size(0), -1, self.head_v_dim)
-        b, _ = self.in_proj_b(hidden_states)
-        a, _ = self.in_proj_a(hidden_states)
+        ba, _ = self.in_proj_ba(hidden_states)
+        b, a = ba.chunk(2, dim=-1)
 
         b = b.contiguous()
         a = a.contiguous()
@@ -506,11 +365,18 @@ class Qwen3_5Model(Qwen3NextModel):
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
+            # self attention
             ("qkv_proj", "q_proj", "q"),
             ("qkv_proj", "k_proj", "k"),
             ("qkv_proj", "v_proj", "v"),
+            # mlp
             ("gate_up_proj", "gate_proj", 0),
             ("gate_up_proj", "up_proj", 1),
+            # GDN
+            ("in_proj_qkvz", "in_proj_qkv", (0, 1, 2)),
+            ("in_proj_qkvz", "in_proj_z", 3),
+            ("in_proj_ba", "in_proj_b", 0),
+            ("in_proj_ba", "in_proj_a", 1),
         ]
 
         params_dict = dict(self.named_parameters())
@@ -657,6 +523,9 @@ class Qwen3_5ForCausalLMBase(
             "v_proj",
         ],
         "gate_up_proj": ["gate_proj", "up_proj"],
+        # GDN fused projections.
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
     }
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -754,7 +623,12 @@ class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLMBase, QwenNextMixtureOfExperts):
     dummy_inputs=Qwen3VLDummyInputsBuilder,
 )
 class Qwen3_5ForConditionalGeneration(Qwen3VLForConditionalGeneration, IsHybrid):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+    packed_modules_mapping = Qwen3VLForConditionalGeneration.packed_modules_mapping | {
+        "in_proj_qkvz": ["in_proj_qkv", "in_proj_z"],
+        "in_proj_ba": ["in_proj_b", "in_proj_a"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
         # protocols have not __init__ method, so we need to use nn.Module.__init__
         nn.Module.__init__(self)
         config: Qwen3_5Config = vllm_config.model_config.hf_config
@@ -962,7 +836,7 @@ class Qwen3_5_MoeMixtureOfExperts(MixtureOfExperts):
 class Qwen3_5MoeForConditionalGeneration(
     Qwen3_5ForConditionalGeneration, Qwen3_5_MoeMixtureOfExperts
 ):
-    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
         # protocols have not __init__ method, so we need to use nn.Module.__init__
         nn.Module.__init__(self)
         config: Qwen3_5MoeConfig = vllm_config.model_config.hf_config

vllm/model_executor/models/qwen3_next.py

@@ -44,6 +44,7 @@ from vllm.model_executor.layers.layernorm import (
 from vllm.model_executor.layers.layernorm import RMSNormGated
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
+    MergedColumnParallelLinear,
     QKVParallelLinear,
     ReplicatedLinear,
     RowParallelLinear,
@@ -406,19 +407,19 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
         self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
 
         # projection of the input hidden states
-        self.projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2
-        self.projection_size_ba = self.num_v_heads * 2
-        self.in_proj_qkvz = ColumnParallelLinear(
-            input_size=self.hidden_size,
-            output_size=self.projection_size_qkvz,
-            bias=False,
+        # Qwen3-Next and Qwen3.5 has a different qkv_proj layout,
+        # we need to create qkvz_proj adaptively here.
+        self.in_proj_qkvz = self.create_qkvz_proj(
+            hidden_size=self.hidden_size,
+            key_dim=self.key_dim,
+            value_dim=self.value_dim,
             quant_config=quant_config,
             prefix=f"{prefix}.in_proj_qkvz",
         )
         # ba_proj doesn't support blockwise fp8 quantization.
-        self.in_proj_ba = ColumnParallelLinear(
+        self.in_proj_ba = MergedColumnParallelLinear(
             input_size=self.hidden_size,
-            output_size=self.projection_size_ba,
+            output_sizes=[self.num_v_heads] * 2,
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.in_proj_ba",
@@ -484,10 +485,26 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
             raise ValueError(f"Duplicate layer name: {prefix}")
         compilation_config.static_forward_context[prefix] = self
 
+    def create_qkvz_proj(
+        self,
+        hidden_size: int,
+        key_dim: int,
+        value_dim: int,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> MergedColumnParallelLinear:
+        return MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[sum((key_dim, key_dim, value_dim, value_dim))],
+            bias=False,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
     def fix_query_key_value_ordering(
         self,
-        mixed_qkvz,
-        mixed_ba,
+        mixed_qkvz: torch.Tensor,
+        mixed_ba: torch.Tensor,
     ):
         """
         Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.

vllm/model_executor/models/siglip.py

@@ -42,9 +42,13 @@ from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalInputs,
     MultiModalKwargsItems,
-    MultiModalUUIDDict,
 )
-from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
+from vllm.multimodal.parse import (
+    ImageProcessorItems,
+    ImageSize,
+    MultiModalDataItems,
+    MultiModalUUIDItems,
+)
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
@@ -189,10 +193,9 @@ class SiglipMultiModalProcessor(BaseMultiModalProcessor[SiglipProcessingInfo]):
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalInputs:
         if mm_items:
             if isinstance(prompt, str):
@@ -221,9 +224,9 @@ class SiglipMultiModalProcessor(BaseMultiModalProcessor[SiglipProcessingInfo]):
         return super().apply(
             prompt=prompt,
             mm_items=mm_items,
+            mm_uuid_items=mm_uuid_items,
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
     def _hf_processor_applies_updates(

vllm/model_executor/models/terratorch.py

@@ -46,7 +46,6 @@ from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalInputs,
     MultiModalKwargsItems,
-    MultiModalUUIDDict,
     PlaceholderRange,
     mm_inputs,
 )
@@ -55,6 +54,7 @@ from vllm.multimodal.parse import (
     ModalityDataItems,
     MultiModalDataItems,
     MultiModalDataParser,
+    MultiModalUUIDItems,
 )
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
@@ -196,15 +196,19 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor[TerratorchProcessing
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalInputs:
+        if hf_processor_mm_kwargs is None:
+            hf_processor_mm_kwargs = {}
         if tokenization_kwargs is None:
             tokenization_kwargs = {}
 
         mm_hashes = self._hash_mm_items(
-            mm_items, hf_processor_mm_kwargs, tokenization_kwargs, mm_uuids=mm_uuids
+            mm_items,
+            mm_uuid_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
         )
 
         _, passthrough_data = self._get_hf_mm_data(mm_items)

vllm/model_executor/models/transformers/multimodal.py

@@ -31,11 +31,14 @@ from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalFieldConfig,
     MultiModalInputs,
-    MultiModalUUIDDict,
     PlaceholderRange,
     mm_inputs,
 )
-from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems
+from vllm.multimodal.parse import (
+    ImageProcessorItems,
+    MultiModalDataItems,
+    MultiModalUUIDItems,
+)
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
@@ -177,9 +180,9 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalInputs:
         """
         Process multi-modal inputs to be used in vLLM.
@@ -187,6 +190,8 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
         Apply HF Processor on prompt text and multi-modal data together,
         outputting token IDs and processed tensors.
         """
+        if hf_processor_mm_kwargs is None:
+            hf_processor_mm_kwargs = {}
         if tokenization_kwargs is None:
             tokenization_kwargs = {}
 
@@ -258,7 +263,9 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
 
         # Use overrides if provided; fallback to data-dependent hashing.
         mm_hashes = self._hash_mm_items(
-            mm_items, hf_processor_mm_kwargs, tokenization_kwargs, mm_uuids=mm_uuids
+            mm_items,
+            mm_uuid_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
         )
 
         return mm_inputs(

vllm/model_executor/models/voxtral.py

@@ -41,13 +41,13 @@ from vllm.multimodal.inputs import (
     MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalKwargsItems,
-    MultiModalUUIDDict,
     NestedTensors,
 )
 from vllm.multimodal.parse import (
     AudioProcessorItems,
     MultiModalDataItems,
     MultiModalDataParser,
+    MultiModalUUIDItems,
 )
 from vllm.multimodal.processing import BaseDummyInputsBuilder, ProcessorInputs
 from vllm.multimodal.processing.processor import (
@@ -363,16 +363,16 @@ class VoxtralMultiModalProcessor(BaseMultiModalProcessor[VoxtralProcessingInfo])
         self,
         prompt: str | list[int],
         mm_data_items: MultiModalDataItems,
+        mm_uuid_items: MultiModalUUIDItems | None,
         hf_processor_mm_kwargs: Mapping[str, object],
         tokenization_kwargs: Mapping[str, object],
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
         prompt_ids, mm_info, _ = super()._cached_apply_hf_processor(
             prompt=prompt,
             mm_data_items=mm_data_items,
+            mm_uuid_items=mm_uuid_items,
             hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
         # NOTE: The tokens are already inserted by the chat template

vllm/multimodal/inputs.py

@@ -155,7 +155,7 @@ The built-in modalities are defined by
 [`MultiModalDataBuiltins`][vllm.multimodal.inputs.MultiModalDataBuiltins].
 """
 
-MultiModalUUIDDict: TypeAlias = Mapping[str, list[str | None] | str]
+MultiModalUUIDDict: TypeAlias = Mapping[str, Sequence[str | None] | str]
 """
 A dictionary containing user-provided UUIDs for items in each modality.
 If a UUID for an item is not provided, its entry will be `None` and

vllm/multimodal/media/connector.py

@@ -146,7 +146,7 @@ class MediaConnector:
 
             connection = self.connection
             data = connection.get_bytes(
-                url,
+                url_spec.url,
                 timeout=fetch_timeout,
                 allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
             )
@@ -177,7 +177,7 @@ class MediaConnector:
 
             connection = self.connection
             data = await connection.async_get_bytes(
-                url,
+                url_spec.url,
                 timeout=fetch_timeout,
                 allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
             )

vllm/multimodal/parse.py

@@ -3,7 +3,7 @@
 
 from abc import ABC, abstractmethod
 from collections import UserDict
-from collections.abc import Callable, Iterator, Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence, Set
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -33,6 +33,7 @@ from .inputs import (
     MultiModalDataDict,
     MultiModalFieldConfig,
     MultiModalKwargsItems,
+    MultiModalUUIDDict,
     VideoItem,
 )
 from .media import MediaWithBytes
@@ -297,14 +298,15 @@ class DictEmbeddingItems(
         return self.data
 
 
-class AudioProcessorItems(ProcessorBatchItems[HfAudioItem]):
-    def __init__(self, data: Sequence[HfAudioItem] | None) -> None:
-        if data is None:
-            data = [None]
+class AudioProcessorItems(ProcessorBatchItems[HfAudioItem | None]):
+    def __init__(self, data: Sequence[HfAudioItem | None]) -> None:
         super().__init__(data, "audio")
 
     def get_audio_length(self, item_idx: int) -> int:
         audio = self.get(item_idx)
+        if audio is None:
+            raise ValueError(f"Cannot get length of cached audio at {item_idx}")
+
         return len(audio)
 
 
@@ -322,14 +324,14 @@ class ImageSize(NamedTuple):
     height: int
 
 
-class ImageProcessorItems(ProcessorBatchItems[HfImageItem]):
-    def __init__(self, data: Sequence[HfImageItem] | None) -> None:
-        if data is None:
-            data = [None]
+class ImageProcessorItems(ProcessorBatchItems[HfImageItem | None]):
+    def __init__(self, data: Sequence[HfImageItem | None]) -> None:
         super().__init__(data, "image")
 
     def get_image_size(self, item_idx: int) -> ImageSize:
         image = self.get(item_idx)
+        if image is None:
+            raise ValueError(f"Cannot get size of cached image at {item_idx}")
 
         if isinstance(image, PILImage.Image):
             return ImageSize(*image.size)
@@ -349,22 +351,31 @@ class ImageEmbeddingItems(EmbeddingItems):
         super().__init__(data, "image", expected_hidden_size)
 
 
-class VideoProcessorItems(ProcessorBatchItems[HfVideoItem]):
+class VideoProcessorItems(ProcessorBatchItems[HfVideoItem | None]):
     def __init__(
         self,
-        data: Sequence[HfVideoItem] | None,
+        data: Sequence[HfVideoItem | None],
         metadata: dict[str, Any] | list[dict[str, Any] | None] | None = None,
     ) -> None:
-        if data is None:
-            data = [None]
         super().__init__(data, "video")
+
         self.metadata = metadata
 
     def get_num_frames(self, item_idx: int) -> int:
-        return len(self.get(item_idx))
+        video = self.get(item_idx)
+        if video is None:
+            raise ValueError(f"Cannot get length of cached video at {item_idx}")
+
+        return len(video)
 
     def get_frame_size(self, item_idx: int) -> ImageSize:
-        image = self.get(item_idx)[0]  # Assume that the video isn't empty
+        video = self.get(item_idx)
+        if video is None:
+            raise ValueError(f"Cannot get size of cached video at {item_idx}")
+        if len(video) == 0:
+            raise ValueError(f"Cannot get size of empty video at {item_idx}")
+
+        image = video[0]
 
         if isinstance(image, PILImage.Image):
             return ImageSize(*image.size)
@@ -400,6 +411,15 @@ class MultiModalDataItems(UserDict[str, ModalityDataItems[Any, Any]]):
     normalized such that each entry corresponds to a list.
     """
 
+    def select(self, modalities: Set[str]):
+        """
+        Construct a new `MultiModalDataItems` instance containing only the
+        selected modalities.
+        """
+        return MultiModalDataItems(
+            {modality: self[modality] for modality in modalities}
+        )
+
     def get_count(self, modality: str, *, strict: bool = True) -> int:
         """
         Get the number of data items belonging to a modality.
@@ -497,19 +517,11 @@ class MultiModalDataParser:
     ) -> TypeGuard[torch.Tensor | list[torch.Tensor]]:
         if isinstance(data, torch.Tensor):
             return data.ndim == 3
-        if is_list_of(data, torch.Tensor):
+        if is_list_of(data, torch.Tensor) and len(data) > 0:
             return data[0].ndim == 2  # type: ignore[index]
 
         return False
 
-    def _is_empty(self, data: object) -> TypeGuard[None]:
-        if isinstance(data, list):
-            return len(data) == 0
-        if isinstance(data, (np.ndarray, torch.Tensor)):
-            return data.size == 0
-
-        return False
-
     def _get_audio_with_sr(
         self,
         audio: AudioItem,
@@ -545,12 +557,6 @@ class MultiModalDataParser:
         data: ModalityData[AudioItem],
     ) -> ModalityDataItems[Any, Any] | None:
         if data is None:
-            return AudioProcessorItems(None)
-
-        # also check single audio item with sampling rate
-        if self._is_empty(data) or (
-            isinstance(data, tuple) and self._is_empty(data[0])
-        ):
             return None
 
         if self.is_embeddings(data):
@@ -558,9 +564,8 @@ class MultiModalDataParser:
 
         data_items: list[AudioItem]
         if (
-            is_list_of(data, float)
-            or isinstance(data, (np.ndarray, torch.Tensor))
-            and data.ndim == 1
+            (is_list_of(data, float) and len(data) > 0)
+            or (isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 1)
             or isinstance(data, tuple)
         ):
             data_items = [data]
@@ -591,18 +596,13 @@ class MultiModalDataParser:
         data: ModalityData[ImageItem],
     ) -> ModalityDataItems[Any, Any] | None:
         if data is None:
-            return ImageProcessorItems(None)
-
-        if self._is_empty(data):
             return None
 
         if self.is_embeddings(data):
             return ImageEmbeddingItems(data, self.expected_hidden_size)
 
-        if (
-            isinstance(data, (PILImage.Image, MediaWithBytes))
-            or isinstance(data, (np.ndarray, torch.Tensor))
-            and data.ndim == 3
+        if isinstance(data, (PILImage.Image, MediaWithBytes)) or (
+            isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 3
         ):
             data_items = [data]
         elif isinstance(data, (np.ndarray, torch.Tensor)):
@@ -617,19 +617,14 @@ class MultiModalDataParser:
         data: ModalityData[VideoItem],
     ) -> ModalityDataItems[Any, Any] | None:
         if data is None:
-            return VideoProcessorItems(None)
-
-        if self._is_empty(data):
             return None
 
         if self.is_embeddings(data):
             return VideoEmbeddingItems(data, self.expected_hidden_size)
 
         data_items: list[VideoItem]
-        if (
-            is_list_of(data, PILImage.Image)
-            or isinstance(data, (np.ndarray, torch.Tensor))
-            and data.ndim == 4
+        if (is_list_of(data, PILImage.Image) and len(data) > 0) or (
+            isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4
         ):
             data_items = [data]
         elif isinstance(data, (np.ndarray, torch.Tensor)):
@@ -664,12 +659,15 @@ class MultiModalDataParser:
         data: ModalityData[Any],
     ) -> ModalityDataItems[Any, Any] | None:
         """Parse vision chunk data (unified image and video chunks)."""
-        if data is None or self._is_empty(data):
+        if data is None:
             return None
+
         if self.is_embeddings(data):
             raise ValueError("Do not support embedding data for vision_chunk right now")
+
         if isinstance(data, dict):
             data = [data]
+
         return VisionChunkProcessorItems(data)
 
     def _get_subparsers(self) -> Mapping[str, ModalityDataParser]:
@@ -693,3 +691,20 @@ class MultiModalDataParser:
                 mm_items[k] = parsed_data
 
         return mm_items
+
+
+MultiModalUUIDItems: TypeAlias = dict[str, Sequence[str | None]]
+"""
+As [`MultiModalUUIDDict`][vllm.multimodal.inputs.MultiModalUUIDDict], but
+normalized such that each entry corresponds to a list.
+"""
+
+
+def parse_mm_uuids(mm_uuids: MultiModalUUIDDict | None) -> MultiModalUUIDItems:
+    if mm_uuids is None:
+        return {}
+
+    return {
+        modality: [uuids] if isinstance(uuids, str) else uuids
+        for modality, uuids in mm_uuids.items()
+    }

vllm/multimodal/processing/processor.py

@@ -32,7 +32,6 @@ from ..inputs import (
     MultiModalKwargsItem,
     MultiModalKwargsItems,
     MultiModalKwargsOptionalItems,
-    MultiModalUUIDDict,
     PlaceholderRange,
     mm_enc_dec_inputs,
     mm_inputs,
@@ -41,6 +40,7 @@ from ..parse import (
     DictEmbeddingItems,
     EmbeddingItems,
     MultiModalDataItems,
+    MultiModalUUIDItems,
 )
 from .context import (
     BaseProcessingInfo,
@@ -1014,11 +1014,15 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         self,
         prompt: str,
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
     ) -> MultiModalInputs:
-        return self.apply(prompt, mm_items, hf_processor_mm_kwargs, mm_uuids=mm_uuids)
+        return self.apply(
+            prompt,
+            mm_items,
+            mm_uuid_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+        )
 
     @abstractmethod
     def _get_mm_fields_config(
@@ -1174,7 +1178,10 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
 
         In addition, return whether prompt updates have been applied.
         """
-        processor_data, passthrough_data = self._get_hf_mm_data(mm_items)
+        valid_mm_items = mm_items.select(
+            {k for k, c in mm_items.get_all_counts().items() if c > 0}
+        )
+        processor_data, passthrough_data = self._get_hf_mm_data(valid_mm_items)
 
         processed_data = self._call_hf_processor(
             prompt=prompt_text,
@@ -1301,69 +1308,57 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
 
     def _hash_mm_items(
         self,
-        mm_items: MultiModalDataItems,
+        mm_data_items: MultiModalDataItems,
+        mm_uuid_items: MultiModalUUIDItems | None,
         hf_processor_mm_kwargs: Mapping[str, object],
-        tokenization_kwargs: Mapping[str, object],
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalHashes:
-        """Create MM hashes to be returned.
-
-
-        Note: When overrides are provided via callers of `apply`,
-        `_hash_mm_items` will be bypassed and the overrides will be used.
-        """
         model_id = self.info.model_id
 
-        hashes: MultiModalHashes = {}
-        mm_uuids = mm_uuids or {}
+        if mm_uuid_items is None:
+            mm_uuid_items = {}
 
-        for modality, items in mm_items.items():
-            if modality in mm_uuids:
-                mm_uuids_per_modality = mm_uuids[modality]
-                if isinstance(mm_uuids_per_modality, str):
-                    mm_uuids_per_modality = [mm_uuids_per_modality]
+        mm_hashes: MultiModalHashes = {}
+        hasher = MultiModalHasher
+
+        for modality, data_items in mm_data_items.items():
+            if modality in mm_uuid_items:
+                uuid_items = mm_uuid_items[modality]
 
                 # For None entries, compute a hash; otherwise, use provided ID.
-                computed: list[str] = []
-                for i, item in enumerate(items.get_all_items_for_hash()):
-                    item_uuid = mm_uuids_per_modality[i]
+                hashes: list[str] = []
+                for i, item in enumerate(data_items.get_all_items_for_hash()):
+                    uuid_item = uuid_items[i]
 
-                    # NOTE: Even if a item_uuid is provided, we still compute a
-                    # hash if `hf_processor_mm_kwargs` or `tokenization_kwargs`
-                    # are provided. This is because the processed multimodal
-                    # inputs can be different depending on the processor kwargs.
-                    if (
-                        item_uuid is None
-                        or hf_processor_mm_kwargs
-                        or tokenization_kwargs
-                    ):
+                    # NOTE: Even if a uuid_item is provided, we still compute a hash
+                    # if `hf_processor_mm_kwargs` is provided.
+                    # This is because the processed multimodal inputs can be different
+                    # depending on the processor kwargs.
+                    if uuid_item is None or hf_processor_mm_kwargs:
                         # NOTE: use provided hash string to hash with kwargs
                         # if available for better performance.
-                        item = item_uuid if item_uuid is not None else item
-                        computed.append(
-                            MultiModalHasher.hash_kwargs(
+                        item = uuid_item if uuid_item is not None else item
+                        hashes.append(
+                            hasher.hash_kwargs(
                                 model_id=model_id,
                                 **{modality: item},
                                 **hf_processor_mm_kwargs,
-                                **tokenization_kwargs,
                             )
                         )
                     else:
-                        computed.append(item_uuid)
-                hashes[modality] = computed
+                        hashes.append(uuid_item)
+
+                mm_hashes[modality] = hashes
             else:
-                hashes[modality] = [
-                    MultiModalHasher.hash_kwargs(
+                mm_hashes[modality] = [
+                    hasher.hash_kwargs(
                         model_id=model_id,
                         **{modality: item},
                         **hf_processor_mm_kwargs,
-                        **tokenization_kwargs,
                     )
-                    for item in items
+                    for item in data_items
                 ]
 
-        return hashes
+        return mm_hashes
 
     def _get_cache_missing_items(
         self,
@@ -1468,10 +1463,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         self,
         prompt: str | list[int],
         mm_data_items: MultiModalDataItems,
+        mm_uuid_items: MultiModalUUIDItems | None,
         hf_processor_mm_kwargs: Mapping[str, object],
         tokenization_kwargs: Mapping[str, object],
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
         (
             prompt_ids,
@@ -1494,9 +1488,8 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         with timed_preprocessor_operation(self.info.ctx, "hashing"):
             mm_hashes = self._hash_mm_items(
                 mm_data_items,
-                hf_processor_mm_kwargs,
-                tokenization_kwargs,
-                mm_uuids=mm_uuids,
+                mm_uuid_items,
+                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             )
 
         mm_prompt_updates = self._get_mm_prompt_updates(
@@ -1517,10 +1510,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         self,
         prompt: str | list[int],
         mm_data_items: MultiModalDataItems,
+        mm_uuid_items: MultiModalUUIDItems | None,
         hf_processor_mm_kwargs: Mapping[str, object],
         tokenization_kwargs: Mapping[str, object],
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> tuple[list[int], MultiModalProcessingInfo, bool]:
         """
         Apply the HF processor on the full prompt text,
@@ -1533,17 +1525,16 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
             return self._apply_hf_processor(
                 prompt=prompt,
                 mm_data_items=mm_data_items,
+                mm_uuid_items=mm_uuid_items,
                 hf_processor_mm_kwargs=hf_processor_mm_kwargs,
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
 
         with timed_preprocessor_operation(self.info.ctx, "hashing"):
             mm_hashes = self._hash_mm_items(
                 mm_data_items,
-                hf_processor_mm_kwargs,
-                tokenization_kwargs,
-                mm_uuids=mm_uuids,
+                mm_uuid_items,
+                hf_processor_mm_kwargs=hf_processor_mm_kwargs,
             )
 
         with timed_preprocessor_operation(self.info.ctx, "cache_lookup"):
@@ -1753,10 +1744,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalInputs:
         """
         Process multi-modal inputs to be used in vLLM.
@@ -1775,6 +1765,8 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         if request_id is not None:
             self.info.ctx.create_timing_stats(request_id)
 
+        if hf_processor_mm_kwargs is None:
+            hf_processor_mm_kwargs = {}
         if tokenization_kwargs is None:
             tokenization_kwargs = {}
 
@@ -1785,9 +1777,9 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
         ) = self._cached_apply_hf_processor(
             prompt,
             mm_items,
+            mm_uuid_items,
             hf_processor_mm_kwargs,
             tokenization_kwargs=tokenization_kwargs,
-            mm_uuids=mm_uuids,
         )
 
         # NOTE: tokenization_kwargs are not required to init processor
@@ -1861,10 +1853,9 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):
         self,
         prompt: str | list[int],
         mm_items: MultiModalDataItems,
-        hf_processor_mm_kwargs: Mapping[str, object],
+        mm_uuid_items: MultiModalUUIDItems | None = None,
+        hf_processor_mm_kwargs: Mapping[str, object] | None = None,
         tokenization_kwargs: Mapping[str, object] | None = None,
-        *,
-        mm_uuids: MultiModalUUIDDict | None = None,
     ) -> MultiModalEncDecInputs:
         """
         Process multi-modal inputs to be used in vLLM.
@@ -1877,9 +1868,9 @@ class EncDecMultiModalProcessor(BaseMultiModalProcessor[_I]):
         encoder_inputs = super().apply(
             encoder_prompt,
             mm_items,
-            hf_processor_mm_kwargs,
-            tokenization_kwargs,
-            mm_uuids=mm_uuids,
+            mm_uuid_items,
+            hf_processor_mm_kwargs=hf_processor_mm_kwargs,
+            tokenization_kwargs=tokenization_kwargs,
         )
 
         return self._get_enc_dec_inputs(

vllm/platforms/interface.py

@@ -17,7 +17,7 @@ if TYPE_CHECKING:
     from torch.distributed import PrefixStore, ProcessGroup
 
     from vllm.config import VllmConfig
-    from vllm.inputs import ProcessorInputs, PromptType
+    from vllm.inputs import ProcessorInputs
     from vllm.pooling_params import PoolingParams
     from vllm.sampling_params import SamplingParams
     from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -568,9 +568,8 @@ class Platform:
     @classmethod
     def validate_request(
         cls,
-        prompt: "PromptType | ProcessorInputs",
-        params: "SamplingParams | PoolingParams",
         processed_inputs: "ProcessorInputs",
+        params: "SamplingParams | PoolingParams",
     ) -> None:
         """Raises if this request is unsupported on this platform"""
 

vllm/reasoning/qwen3_reasoning_parser.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Sequence
+
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
 )
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
 from vllm.entrypoints.openai.responses.protocol import (
     ResponsesRequest,
 )
@@ -12,13 +15,22 @@ from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser
 
 class Qwen3ReasoningParser(BaseThinkingReasoningParser):
     """
-    Reasoning parser for the Qwen3 model.
+    Reasoning parser for the Qwen3/Qwen3.5 model family.
 
-    The Qwen3 model uses <think>...</think> tokens to denote reasoning text
-    within its output. The model provides a strict switch to disable reasoning
-    output via the 'enable_thinking=False' parameter. This parser extracts the
-    reasoning content enclosed by <think> and </think> tokens from the model's
-    output.
+    The Qwen3 model family uses <think>...</think> tokens to denote reasoning
+    text. Starting with Qwen3.5, the chat template places <think> in the
+    prompt so only </think> appears in the generated output. The model
+    provides a strict switch to disable reasoning output via the
+    'enable_thinking=False' parameter.
+
+    When thinking is disabled, the template places <think>\\n\\n</think>\\n\\n
+    in the prompt. The serving layer detects this via prompt_is_reasoning_end
+    and routes deltas as content without calling the streaming parser.
+
+    NOTE: Models up to the 2507 release (e.g., Qwen/Qwen3-235B-A22B-Instruct-2507)
+    use an older chat template where the model generates <think> itself.
+    This parser handles both styles: if <think> appears in the generated output
+    it is stripped before extraction (non-streaming) or skipped (streaming).
     """
 
     @property
@@ -37,31 +49,27 @@ class Qwen3ReasoningParser(BaseThinkingReasoningParser):
         """
         Extract reasoning content from the model output.
 
-        Qwen3 has stricter requirements - it needs both start and end tokens
-        to be present, unlike other models that work with just the end token.
+        The <think> token is placed in the prompt by the chat template,
+        so typically only </think> appears in the generated output.
+        If <think> is present (e.g. from a different template), it is
+        stripped before extraction.
 
-        For text <think>abc</think>xyz:
-        - 'abc' goes to reasoning
-        - 'xyz' goes to content
+        When thinking is disabled (no </think> in output), returns
+        (None, model_output) to indicate all output is content.
 
         Returns:
             tuple[Optional[str], Optional[str]]: reasoning content and content
         """
 
-        # Check if the model output contains both <think> and </think> tokens.
-        if self.start_token not in model_output or self.end_token not in model_output:
-            return None, model_output
-
-        # Check if the <think> is present in the model output, remove it
-        # if it is present.
+        # Strip <think> if present in the generated output.
         model_output_parts = model_output.partition(self.start_token)
         model_output = (
             model_output_parts[2] if model_output_parts[1] else model_output_parts[0]
         )
 
-        # Check if the model output contains the </think> tokens.
-        # If the end token is not found, return the model output as is.
         if self.end_token not in model_output:
+            # No end token means thinking is disabled or the model
+            # did not produce reasoning. Treat everything as content.
             return None, model_output
 
         # Extract reasoning content from the model output.
@@ -69,3 +77,57 @@ class Qwen3ReasoningParser(BaseThinkingReasoningParser):
 
         final_content = content or None
         return reasoning, final_content
+
+    def extract_reasoning_streaming(
+        self,
+        previous_text: str,
+        current_text: str,
+        delta_text: str,
+        previous_token_ids: Sequence[int],
+        current_token_ids: Sequence[int],
+        delta_token_ids: Sequence[int],
+    ) -> DeltaMessage | None:
+        """
+        Extract reasoning content from a streaming delta.
+
+        Since <think> is placed in the prompt by the chat template, all
+        generated tokens before </think> are reasoning and tokens after
+        are content.
+
+        NOTE: When thinking is disabled, no think tokens appear in the
+        generated output. The serving layer detects this via
+        prompt_is_reasoning_end and routes deltas as content without
+        calling this method.
+        """
+        # Strip <think> from delta if present (old template / edge case
+        # where the model generates <think> itself).
+        if self.start_token_id in delta_token_ids:
+            start_idx = delta_text.find(self.start_token)
+            if start_idx >= 0:
+                delta_text = delta_text[start_idx + len(self.start_token) :]
+
+        if self.end_token_id in delta_token_ids:
+            # End token in this delta: split reasoning from content.
+            end_index = delta_text.find(self.end_token)
+            if end_index >= 0:
+                reasoning = delta_text[:end_index]
+                content = delta_text[end_index + len(self.end_token) :]
+                if not reasoning and not content:
+                    return None
+                return DeltaMessage(
+                    reasoning=reasoning if reasoning else None,
+                    content=content if content else None,
+                )
+            # end_token_id in IDs but not in text (already stripped)
+            return None
+
+        # No end token in this delta.
+        if not delta_text:
+            # Nothing left after stripping start token.
+            return None
+        elif self.end_token_id in previous_token_ids:
+            # End token already passed: everything is content now.
+            return DeltaMessage(content=delta_text)
+        else:
+            # No end token yet: still in reasoning phase.
+            return DeltaMessage(reasoning=delta_text)

vllm/renderers/base.py

@@ -51,7 +51,7 @@ if TYPE_CHECKING:
         MultiModalInputs,
         MultiModalUUIDDict,
     )
-    from vllm.multimodal.parse import MultiModalDataItems
+    from vllm.multimodal.parse import MultiModalDataItems, MultiModalUUIDItems
     from vllm.multimodal.processing import BaseMultiModalProcessor
 
 logger = init_logger(__name__)
@@ -463,23 +463,25 @@ class BaseRenderer(ABC, Generic[_T]):
     def _validate_mm_uuids(
         self,
         mm_data: "MultiModalDataDict",
-        mm_items: "MultiModalDataItems",
-        mm_uuids: "MultiModalUUIDDict | None",
+        mm_data_items: "MultiModalDataItems",
+        mm_uuid_items: "MultiModalUUIDItems",
     ) -> None:
-        if mm_uuids is None:
-            mm_uuids = {}
-
-        # NOTE: Keys corresponding to `None` in `mm_data` don't appear in `mm_items`
-        modalities = mm_data.keys() | mm_uuids.keys()
+        # NOTE: Keys corresponding to `None` in `mm_data` don't appear in
+        # `mm_data_items`
+        modalities = mm_data.keys() | mm_uuid_items.keys()
 
         for modality in modalities:
-            data_items = mm_items.get(modality) or list[Any]()
+            data_items = mm_data_items.get(modality)
+            uuid_items = mm_uuid_items.get(modality)
 
-            uuid_items = mm_uuids.get(modality) or list[str | None]()
-            if isinstance(uuid_items, str):
-                uuid_items = [uuid_items]
+            if data_items is None:
+                if uuid_items is None:
+                    raise ValueError(
+                        f"multi_modal_data[{modality!r}] is empty but "
+                        f"multi_modal_uuids[{modality!r}] is missing."
+                    )
 
-            if len(data_items) > 0:
+            elif uuid_items is not None:
                 if len(uuid_items) > 0 and len(data_items) != len(uuid_items):
                     raise ValueError(
                         f"If given, multi_modal_uuids[{modality!r}] must have "
@@ -488,24 +490,17 @@ class BaseRenderer(ABC, Generic[_T]):
                     )
 
                 for i, item in enumerate(data_items):
-                    if item is None:
-                        if not uuid_items:
-                            raise ValueError(
-                                f"multi_modal_data[{modality!r}][{i}] is empty but "
-                                f"multi_modal_uuids[{modality!r}] is missing."
-                            )
-
-                        if uuid_items[i] is None:
-                            raise ValueError(
-                                f"multi_modal_data[{modality!r}][{i}] is empty but "
-                                f"multi_modal_uuids[{modality!r}][{i}] is missing."
-                            )
+                    if item is None and uuid_items[i] is None:
+                        raise ValueError(
+                            f"multi_modal_data[{modality!r}][{i}] is empty but "
+                            f"multi_modal_uuids[{modality!r}][{i}] is missing."
+                        )
 
     def _process_mm_uuids(
         self,
         mm_data: "MultiModalDataDict",
-        mm_items: "MultiModalDataItems",
-        mm_uuids: "MultiModalUUIDDict | None",
+        mm_data_items: "MultiModalDataItems",
+        mm_uuid_items: "MultiModalUUIDItems",
         mm_req_id: str,
     ):
         model_config = self.model_config
@@ -520,40 +515,45 @@ class BaseRenderer(ABC, Generic[_T]):
             and model_config.multimodal_config.mm_processor_cache_gb == 0
             and not self.config.cache_config.enable_prefix_caching
         ):
-            mm_uuids = {
+            mm_uuid_items = {
                 modality: [f"{mm_req_id}-{modality}-{i}" for i in range(data_count)]
-                for modality, data_count in mm_items.get_all_counts().items()
+                for modality, data_count in mm_data_items.get_all_counts().items()
             }
 
-        self._validate_mm_uuids(mm_data, mm_items, mm_uuids)
+        self._validate_mm_uuids(mm_data, mm_data_items, mm_uuid_items)
 
-        return mm_uuids
+        return mm_uuid_items
 
     # TODO: Remove str and tokenization_kwargs after deprecating InputPreprocessor
     def _process_multimodal(
         self,
         prompt: list[int] | str,
         mm_data: "MultiModalDataDict",
+        mm_uuids: "MultiModalUUIDDict | None",
         mm_processor_kwargs: Mapping[str, object] | None,
         tokenization_kwargs: dict[str, Any] | None,
-        mm_uuids: "MultiModalUUIDDict | None",
     ) -> "MultiModalInputs":
+        from vllm.multimodal.parse import parse_mm_uuids
         from vllm.multimodal.processing.context import set_request_id
 
         mm_req_id = f"renderer-mm-{self._mm_req_counter.inc(1)}"
 
         mm_processor = self.get_mm_processor()
 
-        mm_items = mm_processor.info.parse_mm_data(mm_data)
-        mm_uuids = self._process_mm_uuids(mm_data, mm_items, mm_uuids, mm_req_id)
+        mm_data_items = mm_processor.info.parse_mm_data(mm_data)
+        mm_uuid_items = parse_mm_uuids(mm_uuids)
+
+        mm_uuids = self._process_mm_uuids(
+            mm_data, mm_data_items, mm_uuid_items, mm_req_id
+        )
 
         with set_request_id(mm_req_id), set_default_torch_num_threads():
             mm_inputs = mm_processor.apply(
                 prompt,
-                mm_items,
-                hf_processor_mm_kwargs=mm_processor_kwargs or {},
+                mm_data_items,
+                mm_uuid_items,
+                hf_processor_mm_kwargs=mm_processor_kwargs,
                 tokenization_kwargs=tokenization_kwargs,
-                mm_uuids=mm_uuids,
             )
 
         self.update_mm_cache_stats()

vllm/utils/deep_gemm.py

@@ -349,7 +349,7 @@ def _align(x: int, y: int) -> int:
 
 
 # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/v2.1.1/csrc/utils/math.hpp#L19
-def get_tma_aligned_size(x: int, element_size: int):
+def get_tma_aligned_size(x: int, element_size: int) -> int:
     return _align(x, 16 // element_size)
 
 

vllm/v1/engine/async_llm.py

@@ -27,7 +27,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
 from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput
 from vllm.plugins.io_processors import get_io_processor
 from vllm.pooling_params import PoolingParams
-from vllm.renderers import merge_kwargs, renderer_from_config
+from vllm.renderers import renderer_from_config
 from vllm.renderers.inputs.preprocess import extract_prompt_components
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.tasks import SupportedTask
@@ -319,21 +319,6 @@ class AsyncLLM(EngineClient):
                 "prompt logprobs"
             )
 
-        if params.truncate_prompt_tokens is not None:
-            params_type = type(params).__name__
-            warnings.warn(
-                f"The `truncate_prompt_tokens` parameter in `{params_type}` "
-                "is deprecated and will be removed in v0.16. "
-                "Please pass it via `tokenization_kwargs` instead.",
-                DeprecationWarning,
-                stacklevel=2,
-            )
-
-            tokenization_kwargs = merge_kwargs(
-                tokenization_kwargs,
-                dict(truncate_prompt_tokens=params.truncate_prompt_tokens),
-            )
-
         if isinstance(prompt, AsyncGenerator):
             if reasoning_ended is not None:
                 raise NotImplementedError
@@ -353,6 +338,12 @@ class AsyncLLM(EngineClient):
 
         # Convert Input --> Request.
         if isinstance(prompt, EngineCoreRequest):
+            logger.warning_once(
+                "Passing EngineCoreRequest to AsyncLLM.generate() and .add_requests() "
+                "is deprecated and will be removed in v0.18. You should instead pass "
+                "the outputs of Renderer.render_cmpl() or Renderer.render_chat()."
+            )
+
             request = prompt
             if request_id != request.request_id:
                 logger.warning_once(

vllm/v1/engine/input_processor.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import time
+import warnings
 from collections.abc import Mapping
 from typing import Any, Literal
 
@@ -28,6 +29,7 @@ from vllm.sampling_params import SamplingParams
 from vllm.tasks import POOLING_TASKS, SupportedTask
 from vllm.tokenizers import TokenizerLike
 from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
+from vllm.utils.func_utils import supports_kw
 from vllm.utils.jsontree import json_iter_leaves
 from vllm.v1.engine import EngineCoreRequest
 
@@ -72,6 +74,33 @@ class InputProcessor:
             mm_registry=mm_registry,
         )
 
+        from vllm.platforms import current_platform
+
+        platform_validate_request = current_platform.validate_request
+        if supports_kw(platform_validate_request, "prompt"):
+            logger.warning_once(
+                "The signature of Platform.validate_request has changed from "
+                "`(cls, prompt, params, processed_inputs) -> None` to "
+                "`(cls, processed_inputs, params) -> None`. The old signature "
+                "will no longer be supported starting from v0.18."
+            )
+
+            orig_validate_request = platform_validate_request
+
+            def compat_validate_request(
+                processed_inputs: ProcessorInputs,
+                params: SamplingParams | PoolingParams,
+            ):
+                return orig_validate_request(
+                    processed_inputs,
+                    params,
+                    processed_inputs,  # type: ignore
+                )  # type: ignore
+
+            platform_validate_request = compat_validate_request
+
+        self._platform_validate_request = platform_validate_request
+
     @property
     def tokenizer(self) -> TokenizerLike | None:
         return self.renderer.tokenizer
@@ -87,6 +116,16 @@ class InputProcessor:
         supported_tasks: tuple[SupportedTask, ...] | None,
     ):
         """Raise `ValueError` if SamplingParams or PoolingParams is not valid."""
+        if params.truncate_prompt_tokens is not None:
+            params_type = type(params).__name__
+            warnings.warn(
+                f"The `truncate_prompt_tokens` parameter in `{params_type}` "
+                "is deprecated and will be removed in v0.17. "
+                "Please pass it via `tokenization_kwargs` instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
         if isinstance(params, SamplingParams):
             params.verify(
                 self.model_config,
@@ -211,11 +250,24 @@ class InputProcessor:
             )
 
         if isinstance(prompt, dict) and "type" in prompt:
+            if tokenization_kwargs:
+                logger.warning_once(
+                    "Passing tokenization_kwargs to InputProcessor is deprecated "
+                    "and will be removed in v0.18. You should instead pass "
+                    "them to Renderer.render_cmpl() or Renderer.render_chat()."
+                )
+
             if arrival_time is None:
                 arrival_time = prompt.get("arrival_time", time.time())  # type: ignore[assignment]
 
             processed_inputs: ProcessorInputs = prompt  # type: ignore[assignment]
         else:
+            logger.warning_once(
+                "Passing raw prompts to InputProcessor is deprecated "
+                "and will be removed in v0.18. You should instead pass "
+                "the outputs of Renderer.render_cmpl() or Renderer.render_chat()."
+            )
+
             if arrival_time is None:
                 arrival_time = time.time()
 
@@ -224,13 +276,7 @@ class InputProcessor:
                 tokenization_kwargs=tokenization_kwargs,
             )
 
-        from vllm.platforms import current_platform
-
-        current_platform.validate_request(
-            prompt=prompt,
-            params=params,
-            processed_inputs=processed_inputs,
-        )
+        self._platform_validate_request(processed_inputs, params)
 
         encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
         self._validate_model_inputs(encoder_inputs, decoder_inputs)

vllm/v1/engine/llm_engine.py

@@ -234,10 +234,16 @@ class LLMEngine:
 
         # Process raw inputs into the request.
         if isinstance(prompt, EngineCoreRequest):
+            logger.warning_once(
+                "Passing EngineCoreRequest to LLMEngine.generate() and .add_requests() "
+                "is deprecated and will be removed in v0.18. You should instead pass "
+                "the outputs of Renderer.render_cmpl() or Renderer.render_chat()."
+            )
+
             request = prompt
             if request_id != request.request_id:
                 logger.warning_once(
-                    "AsyncLLM.add_request() was passed a request_id parameter that "
+                    "LLMEngine.add_request() was passed a request_id parameter that "
                     "does not match the EngineCoreRequest.request_id attribute. The "
                     "latter will be used, and the former will be ignored."
                 )

vllm/v1/worker/gpu/attn_utils.py

@@ -12,7 +12,6 @@ from vllm.v1.attention.backend import (
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
 )
-from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
     KVCacheConfig,
@@ -144,28 +143,6 @@ def build_slot_mappings_by_layer(
     return slot_mappings_by_layer
 
 
-def prepare_dcp_local_seq_lens(
-    dcp_local_seq_lens: torch.Tensor,
-    seq_lens: torch.Tensor,
-    num_reqs: int,
-    dcp_size: int,
-    dcp_rank: int,
-    cp_kv_cache_interleave_size: int,
-) -> None:
-    """Populate the persistent DCP local seq_lens buffer (CUDA graph safe)."""
-    if dcp_size <= 1:
-        return
-
-    local_seq_lens = get_dcp_local_seq_lens(
-        seq_lens[:num_reqs],
-        dcp_size=dcp_size,
-        dcp_rank=dcp_rank,
-        cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
-    )
-    dcp_local_seq_lens[:num_reqs].copy_(local_seq_lens, non_blocking=True)
-    dcp_local_seq_lens[num_reqs:].zero_()
-
-
 def build_attn_metadata(
     attn_metadata_builders: list[AttentionMetadataBuilder],
     num_reqs: int,
@@ -181,7 +158,6 @@ def build_attn_metadata(
     dcp_local_seq_lens: torch.Tensor | None = None,
 ) -> dict[str, Any]:
     seq_lens = seq_lens[:num_reqs]
-
     if dcp_local_seq_lens is not None:
         dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs]
 

vllm/v1/worker/gpu/block_table.py

@@ -4,7 +4,6 @@ from collections.abc import Iterable
 
 import torch
 
-from vllm.distributed import get_dcp_group
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -19,36 +18,29 @@ class BlockTables:
         max_num_batched_tokens: int,
         max_model_len: int,
         device: torch.device,
-        cp_kv_cache_interleave_size: int = 1,
+        cp_size: int = 1,
+        cp_rank: int = 0,
+        cp_interleave: int = 1,
     ):
         self.block_sizes = block_sizes
         self.max_num_reqs = max_num_reqs
         self.max_num_batched_tokens = max_num_batched_tokens
         self.max_model_len = max_model_len
         self.device = device
-        assert cp_kv_cache_interleave_size >= 1
-        self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
 
-        try:
-            dcp = get_dcp_group()
-            self.dcp_world_size, self.dcp_rank = dcp.world_size, dcp.rank_in_group
-        except AssertionError:
-            self.dcp_world_size, self.dcp_rank = 1, 0
-        # TODO(wentao): PCP supprot
-        self.total_cp_world_size = self.dcp_world_size
-        self.total_cp_rank = self.dcp_rank
+        self.cp_size = cp_size
+        self.cp_rank = cp_rank
+        self.cp_interleave = cp_interleave
 
         self.num_kv_cache_groups = len(self.block_sizes)
         # num_kv_cache_groups x [max_num_reqs, max_num_blocks]
         self.block_tables: list[StagedWriteTensor] = []
         for i in range(self.num_kv_cache_groups):
             block_size = self.block_sizes[i]
-            # with DCP, a request's KV is sharded across
-            # ranks, so one physical block on this rank
-            # corresponds to `block_size * total_cp_world_size`
-            # tokens in the global (unsharded) sequence.
-            virtual_block_size = block_size * self.total_cp_world_size
-            max_num_blocks = cdiv(self.max_model_len, virtual_block_size)
+            # When using DCP, each request's KV cache is sharded among different ranks.
+            # As a result, one block on the current rank covers `block_size * cp_size`
+            # tokens in the full, global (unsharded) sequence.
+            max_num_blocks = cdiv(self.max_model_len, block_size * self.cp_size)
             block_table = StagedWriteTensor(
                 (self.max_num_reqs, max_num_blocks),
                 dtype=torch.int32,
@@ -149,9 +141,9 @@ class BlockTables:
             self.block_sizes_tensor,
             self.slot_mappings,
             self.slot_mappings.stride(0),
-            TOTAL_CP_WORLD_SIZE=self.total_cp_world_size,
-            TOTAL_CP_RANK=self.total_cp_rank,
-            CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,
+            self.cp_rank,
+            CP_SIZE=self.cp_size,
+            CP_INTERLEAVE=self.cp_interleave,
             PAD_ID=PAD_SLOT_ID,
             TRITON_BLOCK_SIZE=1024,  # type: ignore
         )
@@ -204,9 +196,9 @@ def _compute_slot_mappings_kernel(
     block_sizes,  # [num_kv_cache_groups]
     slot_mappings_ptr,  # [num_kv_cache_groups, max_num_tokens]
     slot_mappings_stride,
-    TOTAL_CP_WORLD_SIZE: tl.constexpr,
-    TOTAL_CP_RANK: tl.constexpr,
-    CP_KV_CACHE_INTERLEAVE_SIZE: tl.constexpr,
+    cp_rank,
+    CP_SIZE: tl.constexpr,
+    CP_INTERLEAVE: tl.constexpr,
     PAD_ID: tl.constexpr,
     TRITON_BLOCK_SIZE: tl.constexpr,
 ):
@@ -225,7 +217,6 @@ def _compute_slot_mappings_kernel(
     block_table_ptr = _load_ptr(block_table_ptrs + group_id, tl.int32)
     block_table_stride = tl.load(block_table_strides + group_id)
     block_size = tl.load(block_sizes + group_id)
-    virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE
 
     req_state_idx = tl.load(idx_mapping + batch_idx)
     start_idx = tl.load(query_start_loc + batch_idx)
@@ -233,26 +224,25 @@ def _compute_slot_mappings_kernel(
     for i in range(start_idx, end_idx, TRITON_BLOCK_SIZE):
         offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
         positions = tl.load(pos + offset, mask=offset < end_idx, other=0)
-        block_indices = positions // virtual_block_size
+
+        block_indices = positions // (block_size * CP_SIZE)
+        block_offsets = positions % (block_size * CP_SIZE)
         block_numbers = tl.load(
             block_table_ptr + req_state_idx * block_table_stride + block_indices
         )
-        virtual_block_offsets = positions - block_indices * virtual_block_size
 
-        # determine whether the token is stored on this CP rank.
-        is_local = (
-            virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE
-        ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK
-        # mapping virture block offsets to local block offsets.
-        local_block_offsets = (
-            virtual_block_offsets // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE)
-        ) * CP_KV_CACHE_INTERLEAVE_SIZE + (
-            virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
-        )
+        if CP_SIZE == 1:
+            # Common case: Context parallelism is not used.
+            slot_ids = block_numbers * block_size + block_offsets
+        else:
+            # Context parallelism is used.
+            is_local = block_offsets // CP_INTERLEAVE % CP_SIZE == cp_rank
+            rounds = block_offsets // (CP_INTERLEAVE * CP_SIZE)
+            remainder = block_offsets % CP_INTERLEAVE
+            local_offsets = rounds * CP_INTERLEAVE + remainder
+            slot_ids = block_numbers * block_size + local_offsets
+            slot_ids = tl.where(is_local, slot_ids, PAD_ID)
 
-        # physical slot index
-        slot_ids = block_numbers * block_size + local_block_offsets
-        slot_ids = tl.where(is_local, slot_ids, PAD_ID)
         tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx)
 
 

vllm/v1/worker/gpu/cp_utils.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+def prepare_dcp_local_seq_lens(
+    dcp_local_seq_lens: torch.Tensor,
+    seq_lens: torch.Tensor,
+    num_reqs: int,
+    dcp_size: int,
+    dcp_rank: int,
+    cp_interleave: int,
+) -> None:
+    """Populate the persistent DCP local seq_lens buffer (CUDA graph safe)."""
+    if dcp_size == 1:
+        return
+
+    max_num_reqs = dcp_local_seq_lens.shape[0]
+    BLOCK_SIZE = 128
+    num_blocks = triton.cdiv(max_num_reqs, BLOCK_SIZE)
+    _dcp_local_seq_lens_kernel[(num_blocks,)](
+        dcp_local_seq_lens,
+        seq_lens,
+        dcp_size,
+        dcp_rank,
+        cp_interleave,
+        num_reqs,
+        max_num_reqs,
+        BLOCK_SIZE,
+    )
+
+
+@triton.jit
+def _dcp_local_seq_lens_kernel(
+    out_ptr,
+    seq_lens_ptr,
+    dcp_size,
+    dcp_rank,
+    cp_interleave,
+    num_reqs,
+    max_num_reqs,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    block = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+
+    seq_lens = tl.load(seq_lens_ptr + block, mask=block < num_reqs)
+
+    # Distribute KV cache among different ranks, in a round-robin manner.
+    rounds = seq_lens // (dcp_size * cp_interleave)
+    remainder = seq_lens % (dcp_size * cp_interleave)
+
+    remainder = tl.maximum(remainder - dcp_rank * cp_interleave, 0)
+    remainder = tl.minimum(remainder, cp_interleave)
+    local_seq_lens = rounds * cp_interleave + remainder
+
+    # For [num_reqs, max_num_reqs), pad with 0
+    local_seq_lens = tl.where(block < num_reqs, local_seq_lens, 0)
+    tl.store(out_ptr + block, local_seq_lens, mask=block < max_num_reqs)

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -10,7 +10,6 @@ from tqdm import tqdm
 
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
-from vllm.distributed import get_dcp_group
 from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
 from vllm.forward_context import set_forward_context
 from vllm.v1.attention.backend import AttentionMetadataBuilder
@@ -18,7 +17,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import (
     build_attn_metadata,
     build_slot_mappings_by_layer,
-    prepare_dcp_local_seq_lens,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
@@ -259,22 +257,8 @@ def prepare_inputs_to_capture(
     input_buffers.seq_lens[:num_reqs] = num_tokens
     input_buffers.seq_lens[num_reqs:] = 0
 
-    try:
-        dcp_group = get_dcp_group()
-        dcp_world_size = dcp_group.world_size
-        dcp_rank = dcp_group.rank_in_group
-    except AssertionError:
-        dcp_world_size = 1
-        dcp_rank = 0
-    if dcp_world_size > 1:
-        prepare_dcp_local_seq_lens(
-            input_buffers.dcp_local_seq_lens,
-            input_buffers.seq_lens,
-            num_reqs,
-            dcp_size=dcp_world_size,
-            dcp_rank=dcp_rank,
-            cp_kv_cache_interleave_size=block_tables.cp_kv_cache_interleave_size,
-        )
+    input_buffers.dcp_local_seq_lens[:num_reqs] = num_tokens
+    input_buffers.dcp_local_seq_lens[num_reqs:] = 0
 
     input_block_tables = [x[:num_reqs] for x in block_tables.input_block_tables]
     slot_mappings = block_tables.slot_mappings[:, :num_tokens]

vllm/v1/worker/gpu/model_runner.py

@@ -33,10 +33,10 @@ from vllm.v1.worker.gpu.attn_utils import (
     get_kv_cache_spec,
     init_attn_backend,
     init_kv_cache,
-    prepare_dcp_local_seq_lens,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
 from vllm.v1.worker.gpu.dp_utils import (
     get_cudagraph_and_dp_padding,
@@ -192,6 +192,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.is_first_pp_rank = True
             self.is_last_pp_rank = True
 
+        # Decode context parallelism.
+        self.dcp_size = self.parallel_config.decode_context_parallel_size
+        self.use_dcp = self.dcp_size > 1
+        self.dcp_rank = get_dcp_group().rank_in_group if self.use_dcp else 0
+        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
+
     def update_max_model_len(self, max_model_len: int) -> None:
         self.max_model_len = max_model_len
         self.req_states.max_model_len = max_model_len
@@ -251,9 +257,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             max_num_batched_tokens=self.max_num_tokens,
             max_model_len=self.max_model_len,
             device=self.device,
-            cp_kv_cache_interleave_size=(
-                self.parallel_config.cp_kv_cache_interleave_size
-            ),
+            cp_size=self.dcp_size,
+            cp_rank=self.dcp_rank,
+            cp_interleave=self.cp_interleave,
         )
 
         self.attn_backends, self.attn_metadata_builders = init_attn_backend(
@@ -614,16 +620,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
         max_query_len = num_scheduled_tokens.max().item()
 
-        # Get prefill tokens.
-        prepare_prefill_inputs(
-            self.input_buffers.input_ids,
-            self.req_states.next_prefill_tokens,
-            idx_mapping,
-            query_start_loc,
-            self.req_states.all_token_ids.gpu,
-            self.req_states.prefill_len.gpu,
-            self.req_states.num_computed_tokens.gpu,
-        )
+        # Get prefill tokens if any.
+        if self.req_states.any_prefills(idx_mapping_np):
+            prepare_prefill_inputs(
+                self.input_buffers.input_ids,
+                self.req_states.next_prefill_tokens,
+                idx_mapping,
+                query_start_loc,
+                self.req_states.all_token_ids.gpu,
+                self.req_states.prefill_len.gpu,
+                self.req_states.num_computed_tokens.gpu,
+            )
 
         # Prepare positions and seq_lens.
         prepare_pos_seq_lens(
@@ -635,18 +642,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         )
         seq_lens = self.input_buffers.seq_lens[:num_reqs]
 
-        dcp_size = self.parallel_config.decode_context_parallel_size
-        if dcp_size > 1:
+        if self.use_dcp:
+            # Prepare dcp local seq_lens.
             prepare_dcp_local_seq_lens(
                 self.input_buffers.dcp_local_seq_lens,
-                seq_lens,
+                self.input_buffers.seq_lens,
                 num_reqs,
-                dcp_size=dcp_size,
-                dcp_rank=get_dcp_group().rank_in_group,
-                cp_kv_cache_interleave_size=(
-                    self.parallel_config.cp_kv_cache_interleave_size
-                ),
+                self.dcp_size,
+                self.dcp_rank,
+                self.cp_interleave,
             )
+        dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs]
 
         # Prepare M-RoPE positions.
         if self.uses_mrope:
@@ -695,7 +701,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             block_tables=block_tables,
             slot_mappings=slot_mappings,
             kv_cache_config=self.kv_cache_config,
-            dcp_local_seq_lens=self.input_buffers.dcp_local_seq_lens,
+            dcp_local_seq_lens=dcp_local_seq_lens,
         )
 
         input_ids = self.input_buffers.input_ids[:num_tokens_after_padding]
@@ -1003,15 +1009,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         hidden_states, input_batch, kv_connector_output = self.execute_model_state
         self.execute_model_state = None  # type: ignore
 
-        # Non-last PP rank: hidden_states is None because this rank produced
-        # IntermediateTensors instead of final hidden states. Receive the
-        # sampled tokens broadcast by the last rank and update local state.
         if not self.is_last_pp_rank:
-            received = pp_receive(
+            # Non-last PP rank: hidden_states is None because this rank produced
+            # IntermediateTensors instead of final hidden states. Receive the
+            # sampled tokens broadcast from the last rank and update local state.
+            sampled, num_sampled, num_rejected = pp_receive(
                 input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
             )
-            assert received is not None
-            sampled, num_sampled, num_rejected = received
             self.postprocess(input_batch, sampled, num_sampled, num_rejected)
             return None
 
@@ -1020,8 +1024,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             hidden_states, input_batch, grammar_output
         )
 
-        # Broadcast to non-last PP ranks (handles spec decode multi-token).
         if self.use_pp:
+            # Broadcast to non-last PP ranks (handles spec decode multi-token).
             pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected)
 
         prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(

vllm/v1/worker/gpu/pp_utils.py

@@ -13,8 +13,7 @@ def pp_broadcast(
     num_rejected: torch.Tensor,
 ) -> None:
     pp = get_pp_group()
-    if not pp.is_last_rank:
-        return
+    assert pp.is_last_rank
 
     assert sampled_token_ids.dtype == torch.int64
     torch.distributed.broadcast(
@@ -27,10 +26,9 @@ def pp_broadcast(
 
 def pp_receive(
     num_reqs: int, max_sample_len: int = 1
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     pp = get_pp_group()
-    if pp.is_last_rank:
-        return None
+    assert not pp.is_last_rank
 
     sampled_tokens = torch.empty(
         num_reqs, max_sample_len, dtype=torch.int64, device=pp.device

vllm/v1/worker/gpu/states.py

@@ -60,10 +60,7 @@ class RequestState:
 
         # Last sampled tokens.
         self.last_sampled_tokens = torch.zeros(
-            self.max_num_reqs,
-            1,
-            dtype=torch.int64,
-            device=device,
+            self.max_num_reqs, 1, dtype=torch.int64, device=device
         )
 
         # Draft tokens.
@@ -118,3 +115,9 @@ class RequestState:
             return
         self.index_to_req_id.pop(req_idx, None)
         self.free_indices.append(req_idx)
+
+    def any_prefills(self, idx_mapping_np: np.ndarray) -> bool:
+        return np.any(
+            self.num_computed_prefill_tokens[idx_mapping_np]
+            < self.prefill_len.np[idx_mapping_np]
+        )