@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1356635d83

[New model][Multimodal] Add DeepSeek-V4-Flash-Vision-Exp support (#54566)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

Isotr0py committed Sep 3, 2026 at 02:09 UTC 1356635d837c4ef002ec98c1a0296e7ff60be3c1
37 files changed +2918 -152
csrc/libtorch_stable/moe/moe_ops.h
+3 -1
@@ -28,7 +28,9 @@ void topk_softplus_sqrt(
28 const std::optional<torch::stable::Tensor>& correction_bias,
29 const std::optional<torch::stable::Tensor>& input_ids,
30 const std::optional<torch::stable::Tensor>& tid2eid,
31 - const std::optional<torch::stable::Tensor>& is_padding);
31 + const std::optional<torch::stable::Tensor>& is_padding,
32 + const std::optional<torch::stable::Tensor>& bias_vl,
33 + int64_t image_sentinel_lo);
34
35 void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
36 std::optional<torch::stable::Tensor> topk_ids,
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+172 -28
@@ -81,19 +81,92 @@ __launch_bounds__(128) __global__
81 int num_experts, float routed_scaling_factor,
82 const HashIndType* input_ids,
83 const HashIndType* tid2eid,
84 - const bool* is_padding) {
84 + const bool* is_padding, const float* bias_vl,
85 + int64_t image_sentinel_lo) {
86 const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
87 const int lane = threadIdx.x % 32;
88 if (warp >= num_rows) return;
89 const int64_t token_id = load_index_as_int64(input_ids, warp);
90 const bool is_pad_row = is_padding != nullptr && is_padding[warp];
91 + // Image tokens carry five consecutive in-vocab sentinel ids starting at
92 + // image_sentinel_lo (0 = disabled); they skip the tid2eid lookup and
93 + // select experts by score + bias_vl instead. Ids above the sentinel block
94 + // are regular special tokens and must not match.
95 + const bool is_image = bias_vl != nullptr && image_sentinel_lo > 0 &&
96 + token_id >= image_sentinel_lo &&
97 + token_id < image_sentinel_lo + 5;
98
99 #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
100 cudaGridDependencySynchronize();
101 #endif
102 int expert = 0;
103 float weight = 0.f;
96 - if (lane < 6 && !is_pad_row) {
104 + if (is_image && !is_pad_row) {
105 + // Top-6 by (sqrtsoftplus(x) + bias_vl) over the full row; the routing
106 + // weights are the raw sqrtsoftplus scores of the selected experts.
107 + constexpr int MAX_PER_LANE = 12; // supports num_experts <= 384
108 + const int per_lane = num_experts / 32;
109 + float vals[MAX_PER_LANE];
110 + float sels[MAX_PER_LANE];
111 + #pragma unroll
112 + for (int i = 0; i < MAX_PER_LANE; ++i) {
113 + if (i < per_lane) {
114 + const int e = lane + 32 * i;
115 + const float x = input[warp * num_experts + e];
116 + float v = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
117 + if (isnan(v)) {
118 + v = 0.f;
119 + }
120 + vals[i] = v;
121 + sels[i] = v + bias_vl[e];
122 + } else {
123 + vals[i] = 0.f;
124 + sels[i] = -INFINITY;
125 + }
126 + }
127 + int slot_expert[6];
128 + float slot_weight[6];
129 + #pragma unroll
130 + for (int j = 0; j < 6; ++j) {
131 + float best = -INFINITY;
132 + int best_e = num_experts;
133 + #pragma unroll
134 + for (int i = 0; i < MAX_PER_LANE; ++i) {
135 + if (sels[i] > best) {
136 + best = sels[i];
137 + best_e = lane + 32 * i;
138 + }
139 + }
140 + // Warp argmax reduce; ties go to the lower expert id.
141 + #pragma unroll
142 + for (int mask = 16; mask > 0; mask >>= 1) {
143 + const float other_best = VLLM_SHFL_XOR_SYNC(best, mask);
144 + const int other_e = VLLM_SHFL_XOR_SYNC(best_e, mask);
145 + if (other_best > best || (other_best == best && other_e < best_e)) {
146 + best = other_best;
147 + best_e = other_e;
148 + }
149 + }
150 + float w = 0.f;
151 + #pragma unroll
152 + for (int i = 0; i < MAX_PER_LANE; ++i) {
153 + if (lane + 32 * i == best_e) {
154 + w = vals[i];
155 + sels[i] = -INFINITY;
156 + }
157 + }
158 + #pragma unroll
159 + for (int mask = 16; mask > 0; mask >>= 1) {
160 + w += VLLM_SHFL_XOR_SYNC(w, mask);
161 + }
162 + slot_expert[j] = best_e;
163 + slot_weight[j] = w;
164 + }
165 + if (lane < 6) {
166 + expert = slot_expert[lane];
167 + weight = slot_weight[lane];
168 + }
169 + } else if (lane < 6 && !is_pad_row) {
170 // only load and calculate for 6 experts
171 expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
172 const float x = input[warp * num_experts + expert];
@@ -127,7 +200,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
200 double routed_scaling_factor,
201 const HashIndType* input_ids,
202 const HashIndType* tid2eid, cudaStream_t stream,
130 - const bool* is_padding) {
203 + const bool* is_padding, const float* bias_vl,
204 + int64_t image_sentinel_lo) {
205 if (num_rows == 0) return;
206 auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
207 cudaLaunchConfig_t config = {};
@@ -141,7 +215,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
215 config.numAttrs = 1;
216 const float scale = static_cast<float>(routed_scaling_factor);
217 cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
144 - num_experts, scale, input_ids, tid2eid, is_padding);
218 + num_experts, scale, input_ids, tid2eid, is_padding,
219 + bias_vl, image_sentinel_lo);
220 }
221 #endif
222
@@ -174,7 +249,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
249 const int start_expert, const int end_expert, const bool renormalize,
250 double routed_scaling_factor, const float* correction_bias,
251 const HashIndType* input_ids, const HashIndType* tid2eid,
177 - const bool* is_padding) {
252 + const bool* is_padding, const float* bias_vl,
253 + int64_t image_sentinel_lo) {
254 static_assert(std::is_same_v<InputType, float> ||
255 std::is_same_v<InputType, __nv_bfloat16> ||
256 std::is_same_v<InputType, __half>,
@@ -241,6 +317,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
317 const bool row_is_active = finished ? !finished[thread_row] : true;
318 const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
319
320 + // Image tokens carry five consecutive in-vocab sentinel ids starting at
321 + // image_sentinel_lo (0 = disabled); they select experts with bias_vl
322 + // instead of correction_bias / the tid2eid hash table. Ids above the
323 + // sentinel block are regular special tokens and must not match.
324 + int64_t token_id = 0;
325 + if (input_ids != nullptr) {
326 + token_id = load_index_as_int64(input_ids, thread_row);
327 + }
328 + const bool use_vl_bias = bias_vl != nullptr && image_sentinel_lo > 0 &&
329 + token_id >= image_sentinel_lo &&
330 + token_id < image_sentinel_lo + 5;
331 + const float* row_bias = use_vl_bias ? bias_vl : correction_bias;
332 +
333 // We finally start setting up the read pointers for each thread. First, each
334 // thread jumps to the start of the row it will read.
335 const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
@@ -323,9 +412,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
412 constexpr float threshold = 20.0f;
413 constexpr float beta = 1.0f;
414
326 - // Hash MoE path: indices are predetermined from lookup table
327 - if constexpr (USE_HASH) {
328 - const int64_t token_id = load_index_as_int64(input_ids, thread_row);
415 + // Hash MoE path: indices are predetermined from lookup table. Image rows
416 + // (sentinel ids) fall through to the score-based top-k path with bias_vl.
417 + const bool hash_row = USE_HASH && !use_vl_bias;
418 + if (hash_row) {
419 const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
420 if (!is_pad_row) {
421 #pragma unroll
@@ -417,13 +507,13 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
507 if (isnan(val)) {
508 val = 0.f;
509 }
420 - if (correction_bias) {
510 + if (row_bias) {
511 const int group_id = ii / ELTS_PER_LDG;
512 const int local_id = ii % ELTS_PER_LDG;
513 const int expert_idx = first_elt_read_by_thread +
514 group_id * THREADS_PER_ROW * ELTS_PER_LDG +
515 local_id;
426 - val = val + correction_bias[expert_idx];
516 + val = val + row_bias[expert_idx];
517 }
518 row_chunk[ii] = val;
519 }
@@ -488,8 +578,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
578 // to global memory. (This will be a single) thread per row of the
579 // input/output matrices.
580 const int idx = k * thread_row + k_idx;
491 - if (correction_bias != nullptr && should_process_row) {
492 - max_val -= correction_bias[expert];
581 + if (row_bias != nullptr && should_process_row) {
582 + max_val -= row_bias[expert];
583 }
584 output[idx] = max_val;
585 indices[idx] =
@@ -576,7 +666,8 @@ void topkGatingSoftplusSqrtLauncherHelper(
666 const int start_expert, const int end_expert, const bool renormalize,
667 double routed_scaling_factor, const float* correction_bias,
668 const bool use_hash, const HashIndType* input_ids,
579 - const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
669 + const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding,
670 + const float* bias_vl, int64_t image_sentinel_lo) {
671 static constexpr int BYTES_PER_LDG =
672 MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
673 using Constants =
@@ -605,12 +696,13 @@ void topkGatingSoftplusSqrtLauncherHelper(
696 cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
697 indices, source_row, k, start_expert, end_expert,
698 renormalize, routed_scaling_factor, correction_bias,
608 - input_ids, tid2eid, is_padding);
699 + input_ids, tid2eid, is_padding, bias_vl,
700 + image_sentinel_lo);
701 #else
702 kernel<<<num_blocks, block_dim, 0, stream>>>(
703 input, finished, output, num_rows, indices, source_row, k, start_expert,
704 end_expert, renormalize, routed_scaling_factor, correction_bias,
613 - input_ids, tid2eid, is_padding);
705 + input_ids, tid2eid, is_padding, bias_vl, image_sentinel_lo);
706 #endif
707 })
708 }
@@ -624,7 +716,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
716 gating_output, nullptr, topk_weights, topk_indices, \
717 token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
718 routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
627 - stream, is_padding);
719 + stream, is_padding, bias_vl, image_sentinel_lo);
720 #else
721 #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
722 if (WARP_SIZE == 64) { \
@@ -633,14 +725,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
725 gating_output, nullptr, topk_weights, topk_indices, \
726 token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
727 routed_scaling_factor, correction_bias, use_hash, input_ids, \
636 - tid2eid, stream, is_padding); \
728 + tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \
729 } else if (WARP_SIZE == 32) { \
730 topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
731 MAX_BYTES>( \
732 gating_output, nullptr, topk_weights, topk_indices, \
733 token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
734 routed_scaling_factor, correction_bias, use_hash, input_ids, \
643 - tid2eid, stream, is_padding); \
735 + tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \
736 } else { \
737 assert(false && \
738 "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -654,14 +746,16 @@ void topkGatingSoftplusSqrtKernelLauncher(
746 const int topk, const bool renormalize, double routed_scaling_factor,
747 const float* correction_bias, const bool use_hash,
748 const HashIndType* input_ids, const HashIndType* tid2eid,
657 - cudaStream_t stream, const bool* is_padding) {
749 + cudaStream_t stream, const bool* is_padding, const float* bias_vl,
750 + int64_t image_sentinel_lo) {
751 #ifndef USE_ROCM
752 if constexpr (std::is_same_v<InputType, float>) {
753 if (use_hash && topk == 6 && renormalize &&
754 (num_experts == 256 || num_experts == 384)) {
755 launchDsv4HashTopk<IndType, HashIndType>(
756 gating_output, topk_weights, topk_indices, num_tokens, num_experts,
664 - routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
757 + routed_scaling_factor, input_ids, tid2eid, stream, is_padding,
758 + bias_vl, image_sentinel_lo);
759 return;
760 }
761 }
@@ -761,12 +855,30 @@ void dispatch_topk_softplus_sqrt_launch(
855 const std::optional<torch::stable::Tensor>& correction_bias,
856 const std::optional<torch::stable::Tensor>& input_ids,
857 const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
764 - const std::optional<torch::stable::Tensor>& is_padding) {
858 + const std::optional<torch::stable::Tensor>& is_padding,
859 + const std::optional<torch::stable::Tensor>& bias_vl,
860 + int64_t image_sentinel_lo) {
861 const float* bias_ptr = nullptr;
862 if (correction_bias.has_value()) {
863 bias_ptr = correction_bias.value().const_data_ptr<float>();
864 }
865
866 + const float* bias_vl_ptr = nullptr;
867 + if (bias_vl.has_value()) {
868 + const torch::stable::Tensor& bias_vl_tensor = bias_vl.value();
869 + STD_TORCH_CHECK(input_ids.has_value(),
870 + "input_ids is required when bias_vl is set");
871 + STD_TORCH_CHECK(
872 + bias_vl_tensor.scalar_type() == torch::headeronly::ScalarType::Float,
873 + "bias_vl tensor must be float32");
874 + STD_TORCH_CHECK(bias_vl_tensor.dim() == 1, "bias_vl tensor must be 1D");
875 + STD_TORCH_CHECK(bias_vl_tensor.size(0) == num_experts,
876 + "bias_vl size mismatch, expected: ", num_experts);
877 + STD_TORCH_CHECK(bias_vl_tensor.is_contiguous(),
878 + "bias_vl tensor must be contiguous");
879 + bias_vl_ptr = bias_vl_tensor.const_data_ptr<float>();
880 + }
881 +
882 auto launch = [&](auto* topk_indices_ptr) {
883 using OutIndType =
884 typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
@@ -800,7 +912,8 @@ void dispatch_topk_softplus_sqrt_launch(
912 topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
913 num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
914 bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
803 - tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
915 + tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr,
916 + bias_vl_ptr, image_sentinel_lo);
917 } else {
918 STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
919 torch::headeronly::ScalarType::Int);
@@ -810,7 +923,35 @@ void dispatch_topk_softplus_sqrt_launch(
923 topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
924 num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
925 bias_ptr, true, input_ids.value().const_data_ptr<int>(),
813 - tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
926 + tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr,
927 + bias_vl_ptr, image_sentinel_lo);
928 + }
929 + } else if (bias_vl_ptr != nullptr) {
930 + // Non-hash rows with bias_vl: image tokens (five sentinel ids starting
931 + // at image_sentinel_lo) select experts by score + bias_vl. input_ids
932 + // drives the per-row branch, so
933 + // dispatch the hash-index type on its dtype.
934 + if (input_ids.value().scalar_type() ==
935 + torch::headeronly::ScalarType::Long) {
936 + vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
937 + int64_t>(
938 + gating_output, topk_weights.mutable_data_ptr<float>(),
939 + topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
940 + num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
941 + bias_ptr, false, input_ids.value().const_data_ptr<int64_t>(),
942 + static_cast<const int64_t*>(nullptr), stream, is_padding_ptr,
943 + bias_vl_ptr, image_sentinel_lo);
944 + } else {
945 + STD_TORCH_CHECK(input_ids.value().scalar_type() ==
946 + torch::headeronly::ScalarType::Int);
947 + vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
948 + int>(
949 + gating_output, topk_weights.mutable_data_ptr<float>(),
950 + topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
951 + num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
952 + bias_ptr, false, input_ids.value().const_data_ptr<int>(),
953 + static_cast<const int*>(nullptr), stream, is_padding_ptr,
954 + bias_vl_ptr, image_sentinel_lo);
955 }
956 } else {
957 vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -818,7 +959,8 @@ void dispatch_topk_softplus_sqrt_launch(
959 topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
960 num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
961 bias_ptr, false, static_cast<const OutIndType*>(nullptr),
821 - static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
962 + static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr,
963 + nullptr, 0);
964 }
965 };
966
@@ -843,7 +985,9 @@ void topk_softplus_sqrt(
985 const std::optional<torch::stable::Tensor>& correction_bias,
986 const std::optional<torch::stable::Tensor>& input_ids,
987 const std::optional<torch::stable::Tensor>& tid2eid,
846 - const std::optional<torch::stable::Tensor>& is_padding) {
988 + const std::optional<torch::stable::Tensor>& is_padding,
989 + const std::optional<torch::stable::Tensor>& bias_vl,
990 + int64_t image_sentinel_lo) {
991 const int num_experts = gating_output.size(-1);
992 const auto num_tokens = gating_output.numel() / num_experts;
993 const int topk = topk_weights.size(-1);
@@ -857,21 +1001,21 @@ void topk_softplus_sqrt(
1001 gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
1002 token_expert_indices, num_tokens, num_experts, topk, renormalize,
1003 routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
860 - is_padding);
1004 + is_padding, bias_vl, image_sentinel_lo);
1005 } else if (gating_output.scalar_type() ==
1006 torch::headeronly::ScalarType::Half) {
1007 dispatch_topk_softplus_sqrt_launch<__half>(
1008 reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
1009 topk_weights, topk_indices, token_expert_indices, num_tokens,
1010 num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
867 - input_ids, tid2eid, stream, is_padding);
1011 + input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo);
1012 } else if (gating_output.scalar_type() ==
1013 torch::headeronly::ScalarType::BFloat16) {
1014 dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
1015 reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
1016 topk_weights, topk_indices, token_expert_indices, num_tokens,
1017 num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
874 - input_ids, tid2eid, stream, is_padding);
1018 + input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo);
1019 } else {
1020 STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
1021 gating_output.scalar_type());
csrc/libtorch_stable/moe/torch_bindings.cpp
+2 -1
@@ -20,7 +20,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
20 "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
21 "token_expert_indices, Tensor gating_output, bool renormalize, float "
22 "routed_scaling_factor, Tensor? "
23 - "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
23 + "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding, "
24 + "Tensor? bias_vl=None, int image_sentinel_lo=0) -> ()");
25
26 // Calculate the result of moe by summing up the partial results
27 // from all selected experts. topk_ids/expert_map are optional and, when
docs/models/supported_models.md
+1
@@ -528,6 +528,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
528 | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
529 | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
530 | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
531 +| `DeepseekV4ForConditionalGeneration` | DeepSeek-V4 (vision) | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-V4-Flash-Vision-Exp` | | ✅︎ |
532 | `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + I<sup>E+</sup> | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ |
533 | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I<sup>+</sup>/ V<sup>+</sup> | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ |
534 | `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + I<sup>E+</sup> | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ |
tests/config/test_model_arch_config.py
+31
@@ -465,3 +465,34 @@ def test_draft_model_arch_config(
465 _assert_model_config_methods(
466 model_config, expected, check_head_size=check_head_size
467 )
468 +
469 +
470 +def test_deepseek_v4_convertor_splits_vision_architecture():
471 + """The vision checkpoint shares model_type/architectures with the
472 + text-only DeepSeek-V4; the convertor routes it to the VL wrapper class
473 + by rewriting hf_config.architectures (model-class resolution reads the
474 + raw attribute)."""
475 + from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config
476 + from vllm.transformers_utils.model_arch_config_convertor import (
477 + DeepseekV4ModelArchConfigConvertor,
478 + )
479 +
480 + text_cfg = DeepseekV4Config(architectures=["DeepseekV4ForCausalLM"])
481 + conv = DeepseekV4ModelArchConfigConvertor(text_cfg, text_cfg)
482 + assert conv.get_architectures() == ["DeepseekV4ForCausalLM"]
483 + assert text_cfg.architectures == ["DeepseekV4ForCausalLM"]
484 +
485 + vl_cfg = DeepseekV4Config(
486 + architectures=["DeepseekV4ForCausalLM"], vision_n_layers=32
487 + )
488 + conv = DeepseekV4ModelArchConfigConvertor(vl_cfg, vl_cfg)
489 + assert conv.get_architectures() == ["DeepseekV4ForConditionalGeneration"]
490 + assert vl_cfg.architectures == ["DeepseekV4ForConditionalGeneration"]
491 +
492 + # Speculative-draft configs (set by SpeculativeConfig.hf_config_override)
493 + # keep their own architecture even with a vision tower in the config.
494 + for draft_arch in ("DSparkDraftModel", "DeepSeekV4MTPModel"):
495 + draft_cfg = DeepseekV4Config(architectures=[draft_arch], vision_n_layers=32)
496 + conv = DeepseekV4ModelArchConfigConvertor(draft_cfg, draft_cfg)
497 + assert conv.get_architectures() == [draft_arch]
498 + assert draft_cfg.architectures == [draft_arch]
tests/kernels/moe/test_topk_softplus_sqrt.py
+257 -5
@@ -26,18 +26,46 @@ def _torch_topk_softplus_sqrt(
26 e_score_correction_bias: torch.Tensor | None = None,
27 input_ids: torch.Tensor | None = None,
28 hash_indices_table: torch.Tensor | None = None,
29 + bias_vl: torch.Tensor | None = None,
30 + image_sentinel_lo: int = 0,
31 ):
32 scores = F.softplus(gating_output.float()).sqrt()
33 original_scores = scores
32 - if e_score_correction_bias is not None:
33 - scores_for_choice = scores + e_score_correction_bias.unsqueeze(0)
34 - else:
35 - scores_for_choice = scores
34
35 if hash_indices_table is not None:
36 assert input_ids is not None
39 - topk_ids = hash_indices_table[input_ids.long()]
37 + if bias_vl is None or image_sentinel_lo == 0:
38 + topk_ids = hash_indices_table[input_ids.long()]
39 + else:
40 + image_mask = (input_ids.long() >= image_sentinel_lo) & (
41 + input_ids.long() < image_sentinel_lo + 5
42 + )
43 + topk_ids = hash_indices_table[input_ids.long()]
44 + vl_ids = (scores + bias_vl).topk(topk, dim=-1).indices
45 + topk_ids = torch.where(
46 + image_mask.unsqueeze(-1), vl_ids.to(topk_ids.dtype), topk_ids
47 + )
48 else:
49 + if bias_vl is not None and image_sentinel_lo > 0:
50 + # Image tokens carry five consecutive in-vocab sentinel ids
51 + # starting at image_sentinel_lo and select experts with bias_vl
52 + # instead of the regular bias.
53 + assert input_ids is not None
54 + assert e_score_correction_bias is not None
55 + image_mask = (
56 + (input_ids.long() >= image_sentinel_lo)
57 + & (input_ids.long() < image_sentinel_lo + 5)
58 + ).unsqueeze(-1)
59 + row_bias = torch.where(
60 + image_mask,
61 + bias_vl.unsqueeze(0),
62 + e_score_correction_bias.unsqueeze(0),
63 + )
64 + scores_for_choice = scores + row_bias
65 + elif e_score_correction_bias is not None:
66 + scores_for_choice = scores + e_score_correction_bias.unsqueeze(0)
67 + else:
68 + scores_for_choice = scores
69 # Match the fused kernel's deterministic tie-break: lower expert ids
70 # win when scores are equal. torch.topk does not guarantee which tied
71 # index it selects at the k-th boundary.
@@ -368,3 +396,227 @@ def test_fused_topk_softplus_sqrt_padding(
396 sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref)
397 sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops)
398 torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
399 +
400 +
401 +def _make_mixed_input_ids(
402 + num_tokens: int, image_sentinel_lo: int, dtype: torch.dtype = torch.long
403 +) -> torch.Tensor:
404 + """Mix of text, image-sentinel, and above-sentinel ids.
405 +
406 + Every third token is an image sentinel id (cycling through the five slots
407 + [image_sentinel_lo, image_sentinel_lo + 5)); every sixth token carries an
408 + id just above the sentinel block, simulating the named special tokens
409 + that follow it — those must route as regular text tokens.
410 + """
411 + input_ids = torch.randint(
412 + 0, image_sentinel_lo, (num_tokens,), dtype=torch.long, device="cuda"
413 + )
414 + pos = torch.arange(num_tokens, device="cuda")
415 + image_rows = pos % 3 == 0
416 + input_ids[image_rows] = image_sentinel_lo + (pos[image_rows] // 3) % 5
417 + above_rows = pos % 6 == 3
418 + input_ids[above_rows] = image_sentinel_lo + 5
419 + return input_ids.to(dtype)
420 +
421 +
422 +def _assert_topk_matches(
423 + topk_weights: torch.Tensor,
424 + topk_ids: torch.Tensor,
425 + topk_weights_ref: torch.Tensor,
426 + topk_ids_ref: torch.Tensor,
427 +) -> None:
428 + # Different kernels may return the topk experts in different orders when
429 + # scores tie; sort by expert id before comparing.
430 + sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1)
431 + sorted_ids, idx_ops = topk_ids.sort(dim=-1)
432 + torch.testing.assert_close(
433 + sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0
434 + )
435 +
436 + sorted_w_ref = topk_weights_ref.gather(1, idx_ref)
437 + sorted_w = topk_weights.gather(1, idx_ops.to(torch.long))
438 + torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
439 +
440 +
441 +@pytest.mark.skipif(
442 + not current_platform.is_cuda_alike(),
443 + reason="This test is skipped on non-CUDA platform.",
444 +)
445 +@pytest.mark.parametrize(
446 + ("num_experts", "topk", "renormalize", "dtype"),
447 + [
448 + (256, 6, True, torch.float32), # Triton dsv4 fast path
449 + (256, 8, True, torch.bfloat16), # CUDA topk kernel
450 + (384, 6, False, torch.half),
451 + (512, 16, True, torch.float32),
452 + ],
453 +)
454 +def test_fused_topk_softplus_sqrt_bias_vl(
455 + num_experts: int,
456 + topk: int,
457 + renormalize: bool,
458 + dtype: torch.dtype,
459 +):
460 + """Image sentinel tokens must select experts with bias_vl.
461 +
462 + Rows with ids just above the sentinel block (image_sentinel_lo + 5)
463 + simulate the named special tokens that follow it and must route through
464 + the regular bias path; the reference comparison covers them.
465 + """
466 + torch.manual_seed(0)
467 + num_tokens = 64
468 + hidden_size = 1024
469 + vocab_size = 128
470 + image_sentinel_lo = vocab_size - 8
471 + hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
472 + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
473 + e_score_correction_bias = torch.randn(
474 + (num_experts,), dtype=torch.float32, device="cuda"
475 + )
476 + bias_vl = torch.randn((num_experts,), dtype=torch.float32, device="cuda")
477 + input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo)
478 + assert (input_ids == image_sentinel_lo + 5).any()
479 +
480 + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
481 + gating_output=gating_output,
482 + topk=topk,
483 + renormalize=renormalize,
484 + routed_scaling_factor=1.5,
485 + e_score_correction_bias=e_score_correction_bias,
486 + input_ids=input_ids,
487 + bias_vl=bias_vl,
488 + image_sentinel_lo=image_sentinel_lo,
489 + )
490 +
491 + topk_weights, topk_ids = fused_topk_bias(
492 + hidden_states=hidden_states,
493 + gating_output=gating_output,
494 + scoring_func="sqrtsoftplus",
495 + e_score_correction_bias=e_score_correction_bias,
496 + topk=topk,
497 + renormalize=renormalize,
498 + routed_scaling_factor=1.5,
499 + input_tokens=input_ids,
500 + bias_vl=bias_vl,
501 + image_sentinel_lo=image_sentinel_lo,
502 + )
503 +
504 + _assert_topk_matches(topk_weights, topk_ids, topk_weights_ref, topk_ids_ref)
505 +
506 +
507 +@pytest.mark.skipif(
508 + not current_platform.is_cuda_alike(),
509 + reason="This test is skipped on non-CUDA platform.",
510 +)
511 +@pytest.mark.parametrize(
512 + ("num_experts", "topk", "renormalize", "dtype", "indices_dtype"),
513 + [
514 + (256, 6, True, torch.float32, torch.long), # specialized hash kernel
515 + (384, 6, True, torch.float32, torch.int32), # specialized hash kernel
516 + (384, 8, False, torch.bfloat16, torch.long), # generic hash kernel
517 + (512, 6, True, torch.float32, torch.long), # generic hash kernel
518 + ],
519 +)
520 +def test_fused_topk_softplus_sqrt_hash_bias_vl(
521 + num_experts: int,
522 + topk: int,
523 + renormalize: bool,
524 + dtype: torch.dtype,
525 + indices_dtype: torch.dtype,
526 +):
527 + """Hash MoE: text rows use tid2eid, image rows use topk(score + bias_vl).
528 +
529 + Sentinel ids are in-vocab; rows with ids above the sentinel block
530 + (image_sentinel_lo + 5) must still come straight from the hash table.
531 + """
532 + torch.manual_seed(0)
533 + num_tokens = 64
534 + hidden_size = 1024
535 + vocab_size = 64
536 + image_sentinel_lo = vocab_size - 8
537 + hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
538 + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
539 + hash_indices_table = torch.stack(
540 + [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
541 + ).to(device="cuda", dtype=indices_dtype)
542 + bias_vl = torch.randn((num_experts,), dtype=torch.float32, device="cuda")
543 + input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo)
544 + image_rows = (input_ids >= image_sentinel_lo) & (input_ids < image_sentinel_lo + 5)
545 + assert image_rows.any() and (~image_rows).any()
546 + assert (input_ids == image_sentinel_lo + 5).any()
547 +
548 + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
549 + gating_output=gating_output,
550 + topk=topk,
551 + renormalize=renormalize,
552 + routed_scaling_factor=2.5,
553 + input_ids=input_ids,
554 + hash_indices_table=hash_indices_table,
555 + bias_vl=bias_vl,
556 + image_sentinel_lo=image_sentinel_lo,
557 + )
558 +
559 + topk_weights, topk_ids = fused_topk_bias(
560 + hidden_states=hidden_states,
561 + gating_output=gating_output,
562 + scoring_func="sqrtsoftplus",
563 + e_score_correction_bias=None,
564 + topk=topk,
565 + renormalize=renormalize,
566 + input_tokens=input_ids,
567 + hash_indices_table=hash_indices_table,
568 + routed_scaling_factor=2.5,
569 + bias_vl=bias_vl,
570 + image_sentinel_lo=image_sentinel_lo,
571 + )
572 +
573 + _assert_topk_matches(topk_weights, topk_ids, topk_weights_ref, topk_ids_ref)
574 +
575 + # Text rows must come straight from the hash table, image rows from
576 + # topk(score + bias_vl).
577 + text_ids = hash_indices_table[input_ids[~image_rows].long()]
578 + torch.testing.assert_close(
579 + topk_ids[~image_rows].to(text_ids.dtype), text_ids, atol=0, rtol=0
580 + )
581 +
582 +
583 +@pytest.mark.skipif(
584 + not current_platform.is_cuda(),
585 + reason="The DeepSeek V4 fast path is CUDA-only.",
586 +)
587 +def test_dsv4_fast_topk_bias_vl():
588 + torch.manual_seed(0)
589 + num_tokens = 33
590 + num_experts = 256
591 + vocab_size = 128
592 + image_sentinel_lo = vocab_size - 8
593 + gating_output = torch.randn(
594 + (num_tokens, num_experts), dtype=torch.float32, device="cuda"
595 + )
596 + correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
597 + bias_vl = torch.randn(num_experts, dtype=torch.float32, device="cuda")
598 + input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo)
599 +
600 + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
601 + gating_output=gating_output,
602 + topk=6,
603 + renormalize=True,
604 + routed_scaling_factor=1.5,
605 + e_score_correction_bias=correction_bias,
606 + input_ids=input_ids,
607 + bias_vl=bias_vl,
608 + image_sentinel_lo=image_sentinel_lo,
609 + )
610 + topk_weights, topk_ids = dsv4_topk(
611 + gating_output,
612 + correction_bias,
613 + torch.int64,
614 + 1.5,
615 + input_ids=input_ids,
616 + bias_vl=bias_vl,
617 + image_sentinel_lo=image_sentinel_lo,
618 + )
619 +
620 + assert topk_ids.dtype == torch.int64
621 + torch.testing.assert_close(topk_ids_ref.to(torch.int64), topk_ids, atol=0, rtol=0)
622 + torch.testing.assert_close(topk_weights_ref, topk_weights, atol=2e-5, rtol=2e-5)
tests/models/multimodal/processing/test_tensor_schema.py
+5
@@ -164,6 +164,11 @@ def test_model_tensor_schema(model_id: str):
164 "Kimi-K2.5's offline inference has issues about vision chunks. Fix later."
165 )
166
167 + if model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" and not (
168 + current_platform.is_cuda()
169 + ):
170 + pytest.skip("Deepseek V4 is only supported on CUDA")
171 +
172 model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
173 model_info.check_available_online(on_fail="skip")
174 model_info.check_transformers_version(
tests/models/registry.py
+3
@@ -833,6 +833,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
833 "DeepseekOCR2ForCausalLM": _HfExamplesInfo(
834 "deepseek-ai/DeepSeek-OCR-2",
835 ),
836 + "DeepseekV4ForConditionalGeneration": _HfExamplesInfo(
837 + "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp",
838 + ),
839 "Dots3NoteForCausalLM": _HfExamplesInfo(
840 "dots-studio/dots3-note-prev",
841 is_available_online=False,
tests/models/test_initialization.py
+6
@@ -132,6 +132,12 @@ def can_initialize(
132 f"capability {capability.major}.{capability.minor}"
133 )
134
135 + if model_arch == "DeepseekV4ForConditionalGeneration":
136 + from vllm.platforms import current_platform
137 +
138 + if not current_platform.is_cuda():
139 + pytest.skip("Deepseek V4 is only supported on CUDA")
140 +
141 with (
142 patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1),
143 monkeypatch.context() as m,
tests/models/test_registry.py
+5
@@ -64,6 +64,11 @@ def test_registry_imports(model_arch):
64 ):
65 pytest.skip("HY V4 is only supported on CUDA")
66
67 + if model_arch == "DeepseekV4ForConditionalGeneration" and (
68 + current_platform.is_cuda()
69 + ):
70 + pytest.skip("Deepseek V4 is only supported on CUDA")
71 +
72 # Ensure all model classes can be imported successfully
73 model_cls = ModelRegistry._try_load_model_cls(model_arch)
74 assert model_cls is not None
tests/tokenizers_/test_deepseek_v4.py
+58
@@ -6,6 +6,7 @@ from pathlib import Path
6 from types import SimpleNamespace
7
8 import pytest
9 +import torch
10
11 from vllm.entrypoints.chat_utils import parse_chat_messages
12 from vllm.renderers.registry import RENDERER_REGISTRY
@@ -427,3 +428,60 @@ def test_deepseek_v4_encode_messages_rejects_invalid_arguments(kwargs):
428
429 with pytest.raises(ValueError):
430 encode_messages([{"role": "user", "content": "Hello"}], **kwargs)
431 +
432 +
433 +def test_deepseek_v4_image_blocks_become_placeholders():
434 + prompt = _tokenizer().apply_chat_template(
435 + [
436 + {
437 + "role": "user",
438 + "content": [
439 + {"type": "text", "text": "first:"},
440 + {"type": "image_url", "image_url": {"url": "file:///a.png"}},
441 + {"type": "text", "text": "second:"},
442 + {"type": "image", "source": {"data": "AAAA"}},
443 + ],
444 + }
445 + ],
446 + tokenize=False,
447 + thinking=False,
448 + )
449 +
450 + assert "<｜User｜>first:<｜deepseek_image｜>second:<｜deepseek_image｜>" in prompt
451 +
452 +
453 +def test_deepseek_v4_image_sentinel_ids_match_tokenizer():
454 + """The borrowed sentinel ids must line up with the reserved
455 + ``<|place_holder_mm_span_XXXX|>`` tokens in the tokenizer."""
456 + from vllm.models.deepseek_v4.common.mm_preprocess import (
457 + IMAGE_SENTINEL_BASE_ID,
458 + IMAGE_SENTINEL_TOKEN_NAMES,
459 + image_sentinel_mask,
460 + validate_image_sentinel_ids,
461 + )
462 +
463 + class FakeTokenizer:
464 + def __init__(self, offset: int = 0) -> None:
465 + self.offset = offset
466 +
467 + def convert_tokens_to_ids(self, token: str) -> int:
468 + return (
469 + IMAGE_SENTINEL_BASE_ID
470 + + self.offset
471 + + (IMAGE_SENTINEL_TOKEN_NAMES.index(token))
472 + )
473 +
474 + validate_image_sentinel_ids(FakeTokenizer()) # no raise
475 + with pytest.raises(ValueError, match="sentinel"):
476 + validate_image_sentinel_ids(FakeTokenizer(offset=1))
477 +
478 + ids = torch.tensor(
479 + [
480 + 1,
481 + IMAGE_SENTINEL_BASE_ID,
482 + IMAGE_SENTINEL_BASE_ID + 4,
483 + IMAGE_SENTINEL_BASE_ID + 5,
484 + 129264,
485 + ]
486 + )
487 + assert image_sentinel_mask(ids).tolist() == [False, True, True, False, False]
tests/v1/attention/test_deepseek_v4_swa_visible.py new
+580
@@ -0,0 +1,580 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for DeepSeek-V4 in-image bidirectional SWA visibility (vision variant).
4 +
5 +The reference semantics are a torch port of the official
6 +``get_image_visible`` / ``get_window_topk_idxs_visible``: inside an image span
7 +[span_start, span_end] (inclusive), a token at ``pos`` sees
8 +``min(pos - span_start, max_image_tokens - 1)`` extra tokens to the left and
9 +``min(span_end - pos, max_image_tokens)`` to the right; the window then starts
10 +at ``max(pos - (window - 1) - max(left - (window - 1), 0), 0)`` and ends at
11 +``pos + right`` (inclusive).
12 +"""
13 +
14 +import pytest
15 +import torch
16 +from typing_extensions import TypedDict
17 +
18 +from tests.v1.attention.utils import create_vllm_config
19 +from vllm.models.deepseek_v4.common.ops.cache_utils import (
20 + build_flashinfer_mixed_sparse_indices,
21 + combine_topk_swa_indices,
22 +)
23 +from vllm.v1.attention.backend import CommonAttentionMetadata
24 +from vllm.v1.attention.backends.mla.sparse_swa import (
25 + DeepseekSparseSWAMetadataBuilder,
26 + _compute_image_visibility_kernel,
27 + _compute_swa_indices_and_lens_kernel,
28 +)
29 +from vllm.v1.kv_cache_interface import SlidingWindowMLASpec
30 +
31 +WINDOW = 8
32 +MAX_IMG = 6
33 +WIDTH = WINDOW + MAX_IMG
34 +BLOCK_SIZE = 64
35 +
36 +
37 +def ref_left_right(
38 + seq_lens: list[int],
39 + query_lens: list[int],
40 + spans_per_req: list[list[tuple[int, int]]],
41 + max_image_tokens: int,
42 +) -> tuple[list[int], list[int]]:
43 + """Per-token (left, right) for the flattened decode-first token stream."""
44 + lefts: list[int] = []
45 + rights: list[int] = []
46 + for seq_len, query_len, spans in zip(seq_lens, query_lens, spans_per_req):
47 + prefix_len = seq_len - query_len
48 + for i in range(query_len):
49 + pos = prefix_len + i
50 + left = right = 0
51 + for span_start, span_end in spans:
52 + if span_start <= pos <= span_end:
53 + left = min(pos - span_start, max_image_tokens - 1)
54 + right = min(span_end - pos, max_image_tokens)
55 + lefts.append(left)
56 + rights.append(right)
57 + return lefts, rights
58 +
59 +
60 +def ref_swa_bounds(pos: int, window: int, left: int, right: int) -> tuple[int, int]:
61 + """Reference [start, end) window bounds for one query token."""
62 + left_add = max(left - (window - 1), 0)
63 + start = max(pos - (window - 1) - left_add, 0)
64 + return start, pos + right + 1
65 +
66 +
67 +def make_batch(
68 + seq_lens: list[int],
69 + query_lens: list[int],
70 + device: torch.device,
71 +):
72 + """Build the kernel inputs for a decode-first batch."""
73 + num_reqs = len(seq_lens)
74 + query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
75 + query_start_loc[1:] = torch.tensor(
76 + query_lens, dtype=torch.int32, device=device
77 + ).cumsum(0)
78 + num_tokens = int(query_start_loc[-1])
79 + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
80 + token_to_req = torch.repeat_interleave(
81 + torch.arange(num_reqs, dtype=torch.int32, device=device),
82 + torch.tensor(query_lens, dtype=torch.int32, device=device),
83 + )
84 + max_blocks = (max(seq_lens) + BLOCK_SIZE - 1) // BLOCK_SIZE
85 + block_table = torch.arange(
86 + num_reqs * max_blocks, dtype=torch.int32, device=device
87 + ).view(num_reqs, max_blocks)
88 + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
89 + return query_start_loc, seq_lens_t, token_to_req, slot_mapping, block_table
90 +
91 +
92 +def ref_swa_slot_rows(
93 + seq_lens: list[int],
94 + query_lens: list[int],
95 + spans_per_req: list[list[tuple[int, int]]],
96 + block_table: torch.Tensor,
97 + window: int,
98 + max_image_tokens: int,
99 + width: int,
100 +) -> tuple[list[list[int]], list[int]]:
101 + """Reference paged slot-id rows and lens for every token in the batch."""
102 + block_table_cpu = block_table.cpu()
103 + lefts, rights = ref_left_right(
104 + seq_lens, query_lens, spans_per_req, max_image_tokens
105 + )
106 + rows: list[list[int]] = []
107 + lens: list[int] = []
108 + token = 0
109 + for req, (seq_len, query_len) in enumerate(zip(seq_lens, query_lens)):
110 + prefix_len = seq_len - query_len
111 + for i in range(query_len):
112 + pos = prefix_len + i
113 + start, end = ref_swa_bounds(pos, window, lefts[token], rights[token])
114 + row = []
115 + for p in range(start, end):
116 + blk = int(block_table_cpu[req, p // BLOCK_SIZE])
117 + row.append(blk * BLOCK_SIZE + p % BLOCK_SIZE)
118 + lens.append(len(row))
119 + row.extend([-1] * (width - len(row)))
120 + rows.append(row)
121 + token += 1
122 + return rows, lens
123 +
124 +
125 +def run_swa_kernel(
126 + seq_lens: list[int],
127 + query_lens: list[int],
128 + spans_per_req: list[list[tuple[int, int]]],
129 + window: int = WINDOW,
130 + max_image_tokens: int = MAX_IMG,
131 + with_image: bool = True,
132 +) -> tuple[torch.Tensor, torch.Tensor]:
133 + device = torch.device("cuda")
134 + query_start_loc, seq_lens_t, token_to_req, slot_mapping, block_table = make_batch(
135 + seq_lens, query_lens, device
136 + )
137 + num_tokens = int(query_start_loc[-1])
138 + width = window + (max_image_tokens if with_image else 0)
139 + swa_indices = torch.zeros(num_tokens, 1, width, dtype=torch.int32, device=device)
140 + swa_lens = torch.zeros(num_tokens, dtype=torch.int32, device=device)
141 + is_valid = slot_mapping >= 0
142 +
143 + if with_image:
144 + lefts, rights = ref_left_right(
145 + seq_lens, query_lens, spans_per_req, max_image_tokens
146 + )
147 + left_t = torch.tensor(lefts, dtype=torch.int32, device=device)
148 + right_t = torch.tensor(rights, dtype=torch.int32, device=device)
149 + else:
150 + left_t = right_t = swa_lens # unused dummies
151 +
152 + _compute_swa_indices_and_lens_kernel[(num_tokens,)](
153 + swa_indices,
154 + swa_indices.stride(0),
155 + swa_lens,
156 + window,
157 + width,
158 + left_t,
159 + right_t,
160 + query_start_loc,
161 + seq_lens_t,
162 + token_to_req,
163 + is_valid,
164 + block_table,
165 + block_table.stride(0),
166 + BLOCK_SIZE,
167 + token_offset=0,
168 + HAS_IMAGE=with_image,
169 + TRITON_BLOCK_SIZE=1024,
170 + )
171 + return swa_indices[:, 0], swa_lens
172 +
173 +
174 +# seq_lens, query_lens, spans per request (positions are prompt-absolute,
175 +# inclusive on both ends, matching the reference's sentinel-bracketed spans).
176 +class _Case(TypedDict):
177 + seq_lens: list[int]
178 + query_lens: list[int]
179 + spans: list[list[tuple[int, int]]]
180 +
181 +
182 +CASES: list[_Case] = [
183 + # two image spans in one request, one pure-text request
184 + {
185 + "seq_lens": [30, 12],
186 + "query_lens": [30, 12],
187 + "spans": [[(4, 12), (20, 24)], []],
188 + },
189 + # span larger than max_image_tokens, straddling the window boundary
190 + {
191 + "seq_lens": [26],
192 + "query_lens": [26],
193 + "spans": [[(2, 17)]],
194 + },
195 + # chunked prefill: span fully inside the second chunk's query
196 + {
197 + "seq_lens": [40, 9],
198 + "query_lens": [16, 9],
199 + "spans": [[(30, 38)], []],
200 + },
201 + # span ending exactly at the prompt end; tiny request
202 + {
203 + "seq_lens": [7, 15],
204 + "query_lens": [7, 15],
205 + "spans": [[(0, 6)], [(10, 14)]],
206 + },
207 +]
208 +
209 +
210 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
211 +@pytest.mark.parametrize("case", CASES)
212 +def test_swa_indices_kernel_with_image_spans(case):
213 + rows, lens = ref_swa_slot_rows(
214 + case["seq_lens"],
215 + case["query_lens"],
216 + case["spans"],
217 + make_batch(case["seq_lens"], case["query_lens"], torch.device("cuda"))[4],
218 + WINDOW,
219 + MAX_IMG,
220 + WIDTH,
221 + )
222 + indices, actual_lens = run_swa_kernel(
223 + case["seq_lens"], case["query_lens"], case["spans"], with_image=True
224 + )
225 + assert actual_lens.cpu().tolist() == lens
226 + assert indices.cpu().tolist() == rows
227 +
228 +
229 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
230 +@pytest.mark.parametrize("case", CASES)
231 +def test_swa_indices_kernel_without_image_unchanged(case):
232 + """HAS_IMAGE=False must reproduce the plain causal sliding window."""
233 + rows, lens = ref_swa_slot_rows(
234 + case["seq_lens"],
235 + case["query_lens"],
236 + [[] for _ in case["seq_lens"]],
237 + make_batch(case["seq_lens"], case["query_lens"], torch.device("cuda"))[4],
238 + WINDOW,
239 + MAX_IMG,
240 + WINDOW,
241 + )
242 + indices, actual_lens = run_swa_kernel(
243 + case["seq_lens"], case["query_lens"], case["spans"], with_image=False
244 + )
245 + assert indices.shape[-1] == WINDOW
246 + assert actual_lens.cpu().tolist() == lens
247 + assert indices.cpu().tolist() == rows
248 +
249 +
250 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
251 +def test_image_visibility_kernel():
252 + """The builder's visibility kernel must match get_image_visible."""
253 + device = torch.device("cuda")
254 + seq_lens = [30, 12]
255 + query_lens = [30, 12]
256 + spans = [[(4, 12), (20, 24)], []]
257 + query_start_loc, seq_lens_t, token_to_req, _, _ = make_batch(
258 + seq_lens, query_lens, device
259 + )
260 + num_tokens = int(query_start_loc[-1])
261 +
262 + # CSR span layout: request -> contiguous [start, end) span list.
263 + indptr = [0, 2, 2]
264 + starts = [4, 20]
265 + ends = [12, 24]
266 + left = torch.zeros(num_tokens, dtype=torch.int32, device=device)
267 + right = torch.zeros(num_tokens, dtype=torch.int32, device=device)
268 + _compute_image_visibility_kernel[(num_tokens,)](
269 + left,
270 + right,
271 + torch.tensor(indptr, dtype=torch.int32, device=device),
272 + torch.tensor(starts, dtype=torch.int32, device=device),
273 + torch.tensor(ends, dtype=torch.int32, device=device),
274 + query_start_loc,
275 + seq_lens_t,
276 + token_to_req,
277 + MAX_IMG,
278 + token_offset=0,
279 + )
280 + ref_left, ref_right = ref_left_right(seq_lens, query_lens, spans, MAX_IMG)
281 + assert left.cpu().tolist() == ref_left
282 + assert right.cpu().tolist() == ref_right
283 +
284 +
285 +def combine_case(
286 + compress_ratio: int,
287 + topk: int,
288 + seq_lens: list[int],
289 + query_lens: list[int],
290 + spans: list[list[tuple[int, int]]],
291 + with_image: bool,
292 +):
293 + """Run combine_topk_swa_indices and return (indices, lens, expected)."""
294 + device = torch.device("cuda")
295 + num_reqs = len(seq_lens)
296 + query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
297 + query_start_loc[1:] = torch.tensor(
298 + query_lens, dtype=torch.int32, device=device
299 + ).cumsum(0)
300 + num_tokens = int(query_start_loc[-1])
301 + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
302 + gather_lens = torch.tensor(
303 + [q + min(s - q, WINDOW - 1) for s, q in zip(seq_lens, query_lens)],
304 + dtype=torch.int32,
305 + device=device,
306 + )
307 + N = (max(seq_lens) + compress_ratio - 1) // compress_ratio
308 + M = N + int(gather_lens.max()) + 8
309 + gen = torch.Generator(device="cpu").manual_seed(0)
310 + topk_indices = torch.randint(
311 + 0, 4096, (num_tokens, max(topk, 1)), generator=gen, dtype=torch.int32
312 + ).to(device)
313 + topk_indices = topk_indices[:, : max(topk, 1)]
314 +
315 + if with_image:
316 + lefts, rights = ref_left_right(seq_lens, query_lens, spans, MAX_IMG)
317 + left_t = torch.tensor(lefts, dtype=torch.int32, device=device)
318 + right_t = torch.tensor(rights, dtype=torch.int32, device=device)
319 + else:
320 + left_t = right_t = None
321 +
322 + combined_indices, combined_lens = combine_topk_swa_indices(
323 + topk_indices,
324 + query_start_loc,
325 + seq_lens_t,
326 + gather_lens,
327 + WINDOW,
328 + compress_ratio,
329 + topk,
330 + M,
331 + N,
332 + left_visible=left_t,
333 + right_visible=right_t,
334 + max_image_tokens=MAX_IMG,
335 + )
336 +
337 + # Reference rows.
338 + lefts, rights = ref_left_right(
339 + seq_lens,
340 + query_lens,
341 + spans if with_image else [[] for _ in seq_lens],
342 + MAX_IMG,
343 + )
344 + topk_cpu = topk_indices.cpu()
345 + width = WINDOW + MAX_IMG
346 + combined_topk = (topk + width + 127) // 128 * 128
347 + rows = []
348 + lens = []
349 + token = 0
350 + for b, (seq_len, query_len) in enumerate(zip(seq_lens, query_lens)):
351 + prefix_len = seq_len - query_len
352 + gather_start = seq_len - int(gather_lens[b])
353 + for i in range(query_len):
354 + pos = prefix_len + i
355 + topk_len = min((pos + 1) // compress_ratio, topk)
356 + start, end = ref_swa_bounds(pos, WINDOW, lefts[token], rights[token])
357 + swa_len = end - start
358 + row = [-1] * combined_topk
359 + for j in range(topk_len):
360 + row[j] = int(topk_cpu[token, j]) + M * b
361 + for j in range(swa_len):
362 + row[topk_len + j] = M * b + N + start + j - gather_start
363 + rows.append(row)
364 + lens.append(topk_len + swa_len)
365 + token += 1
366 + return combined_indices, combined_lens, rows, lens
367 +
368 +
369 +COMBINE_CASES = [
370 + dict(compress_ratio=1, topk=0), # SWA-only layer
371 + dict(compress_ratio=4, topk=16), # C4A layer
372 +]
373 +
374 +
375 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
376 +@pytest.mark.parametrize("cfg", COMBINE_CASES)
377 +def test_combine_topk_swa_with_image_spans(cfg):
378 + case = CASES[0]
379 + indices, lens, rows, exp_lens = combine_case(
380 + cfg["compress_ratio"],
381 + cfg["topk"],
382 + case["seq_lens"],
383 + case["query_lens"],
384 + case["spans"],
385 + with_image=True,
386 + )
387 + assert lens.cpu().tolist() == exp_lens
388 + assert indices.cpu().tolist() == rows
389 +
390 +
391 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
392 +@pytest.mark.parametrize("cfg", COMBINE_CASES)
393 +def test_combine_topk_swa_without_image_unchanged(cfg):
394 + """left_visible=None must reproduce the plain causal combined indices."""
395 + case = CASES[0]
396 + indices, lens, rows, exp_lens = combine_case(
397 + cfg["compress_ratio"],
398 + cfg["topk"],
399 + case["seq_lens"],
400 + case["query_lens"],
401 + case["spans"],
402 + with_image=False,
403 + )
404 + assert lens.cpu().tolist() == exp_lens
405 + assert indices.cpu().tolist() == rows
406 +
407 +
408 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
409 +def test_flashinfer_mixed_sparse_indices_with_image_spans():
410 + device = torch.device("cuda")
411 + # 1 decode token (req 0) + two prefill requests (reqs 1, 2).
412 + seq_lens = [20, 30, 9]
413 + query_lens = [1, 30, 9]
414 + spans = [[], [(4, 12)], [(0, 6)]]
415 + query_start_loc = torch.tensor([0, 1, 31, 40], dtype=torch.int32, device=device)
416 + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
417 + token_to_req = torch.tensor(
418 + [0] + [1] * 30 + [2] * 9, dtype=torch.int32, device=device
419 + )
420 + block_table = torch.arange(3, dtype=torch.int32, device=device).view(3, 1)
421 + decode_swa = torch.tensor(
422 + [[100, 101, 102, 103, 104, 105, 106, 107]], dtype=torch.int32, device=device
423 + )
424 + prefill_topk = torch.zeros(39, 0, dtype=torch.int32, device=device)
425 +
426 + lefts, rights = ref_left_right(seq_lens, query_lens, spans, MAX_IMG)
427 + left_t = torch.tensor(lefts, dtype=torch.int32, device=device)
428 + right_t = torch.tensor(rights, dtype=torch.int32, device=device)
429 +
430 + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices(
431 + decode_swa_indices=decode_swa,
432 + decode_compressed_indices=None,
433 + decode_compressed_topk_lens=None,
434 + prefill_topk_indices=prefill_topk,
435 + query_start_loc=query_start_loc,
436 + seq_lens=seq_lens_t,
437 + token_to_req_indices=token_to_req,
438 + swa_block_table=block_table,
439 + swa_block_size=BLOCK_SIZE,
440 + compressed_block_table=None,
441 + compressed_block_size=BLOCK_SIZE,
442 + window_size=WINDOW,
443 + compress_ratio=1,
444 + topk=0,
445 + prefill_left_visible=left_t,
446 + prefill_right_visible=right_t,
447 + max_image_tokens=MAX_IMG,
448 + )
449 +
450 + swa_total = WINDOW + MAX_IMG
451 + assert sparse_indices.shape == (40, swa_total)
452 + assert sparse_lens.cpu().tolist() == [swa_total] * 40
453 + # Decode row: slots copied, image-extension columns padded with -1.
454 + decode_row = sparse_indices[0].cpu().tolist()
455 + assert decode_row[:WINDOW] == list(range(100, 108))
456 + assert decode_row[WINDOW:] == [-1] * MAX_IMG
457 + # Prefill rows: paged slot ids over the widened window.
458 + exp_rows, _ = ref_swa_slot_rows(
459 + seq_lens, query_lens, spans, block_table, WINDOW, MAX_IMG, swa_total
460 + )
461 + actual = sparse_indices[1:].cpu().tolist()
462 + assert actual == exp_rows[1:]
463 +
464 +
465 +def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder:
466 + overrides: dict = {"sliding_window": WINDOW}
467 + if vision:
468 + overrides.update(vision_n_layers=2, vision_max_n_token=MAX_IMG)
469 + vllm_config = create_vllm_config(
470 + max_model_len=4096,
471 + max_num_batched_tokens=64,
472 + max_num_seqs=8,
473 + hf_config_override=overrides,
474 + )
475 + spec = SlidingWindowMLASpec(
476 + block_size=BLOCK_SIZE,
477 + num_kv_heads=1,
478 + head_size=512,
479 + dtype=torch.bfloat16,
480 + sliding_window=WINDOW,
481 + cache_dtype_str="auto",
482 + model_version="deepseek_v4",
483 + )
484 + return DeepseekSparseSWAMetadataBuilder(
485 + kv_cache_spec=spec,
486 + layer_names=["layer0"],
487 + vllm_config=vllm_config,
488 + device=torch.device("cuda"),
489 + )
490 +
491 +
492 +def build_metadata(
493 + builder: DeepseekSparseSWAMetadataBuilder,
494 + seq_lens: list[int],
495 + query_lens: list[int],
496 + mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None,
497 +):
498 + device = torch.device("cuda")
499 + query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch(
500 + seq_lens, query_lens, device
501 + )
502 + return builder.build(
503 + 0,
504 + CommonAttentionMetadata(
505 + query_start_loc=query_start_loc,
506 + query_start_loc_cpu=query_start_loc.cpu(),
507 + seq_lens=seq_lens_t,
508 + seq_lens_cpu_upper_bound=seq_lens_t.cpu(),
509 + num_reqs=len(seq_lens),
510 + num_actual_tokens=int(query_start_loc[-1]),
511 + max_query_len=max(query_lens),
512 + max_seq_len=max(seq_lens),
513 + block_table_tensor=block_table,
514 + slot_mapping=slot_mapping,
515 + causal=True,
516 + mm_req_doc_ranges=mm_req_doc_ranges,
517 + ),
518 + )
519 +
520 +
521 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
522 +def test_builder_in_image_visibility():
523 + seq_lens = [30, 12]
524 + query_lens = [30, 12]
525 + spans = [[(4, 12), (20, 24)], []]
526 + builder = make_builder(vision=True)
527 + md = build_metadata(builder, seq_lens, query_lens, {0: spans[0], 1: spans[1]})
528 + assert md.num_prefills == 2
529 + assert md.num_decode_tokens == 0
530 + assert md.prefill_swa_indices.shape[-1] == WIDTH
531 + assert md.prefill_left_visible is not None
532 + assert md.prefill_right_visible is not None
533 +
534 + ref_left, ref_right = ref_left_right(seq_lens, query_lens, spans, MAX_IMG)
535 + assert md.prefill_left_visible.cpu().tolist() == ref_left
536 + assert md.prefill_right_visible.cpu().tolist() == ref_right
537 +
538 + _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda"))
539 + rows, lens = ref_swa_slot_rows(
540 + seq_lens, query_lens, spans, block_table, WINDOW, MAX_IMG, WIDTH
541 + )
542 + assert md.prefill_swa_lens.cpu().tolist() == lens
543 + assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows
544 +
545 +
546 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
547 +def test_builder_no_image_spans_fast_path():
548 + """Vision model, image-free batch: no visibility tensors, plain window."""
549 + seq_lens = [30, 12]
550 + query_lens = [30, 12]
551 + builder = make_builder(vision=True)
552 + md = build_metadata(builder, seq_lens, query_lens, {0: [], 1: []})
553 + assert md.prefill_left_visible is None
554 + assert md.prefill_right_visible is None
555 +
556 + _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda"))
557 + rows, lens = ref_swa_slot_rows(
558 + seq_lens, query_lens, [[], []], block_table, WINDOW, MAX_IMG, WIDTH
559 + )
560 + assert md.prefill_swa_lens.cpu().tolist() == lens
561 + assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows
562 +
563 +
564 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
565 +def test_builder_text_model_unchanged():
566 + """Text-only model: buffers stay window-sized and spans are ignored."""
567 + seq_lens = [30, 12]
568 + query_lens = [30, 12]
569 + builder = make_builder(vision=False)
570 + assert builder.max_image_tokens == 0
571 + md = build_metadata(builder, seq_lens, query_lens, None)
572 + assert md.prefill_swa_indices.shape[-1] == WINDOW
573 + assert md.prefill_left_visible is None
574 +
575 + _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda"))
576 + rows, lens = ref_swa_slot_rows(
577 + seq_lens, query_lens, [[], []], block_table, WINDOW, MAX_IMG, WINDOW
578 + )
579 + assert md.prefill_swa_lens.cpu().tolist() == lens
580 + assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows
vllm/_custom_ops.py
+4
@@ -2501,6 +2501,8 @@ def topk_hash_softplus_sqrt(
2501 input_tokens: torch.Tensor | None = None,
2502 hash_indices_table: torch.Tensor | None = None,
2503 is_padding: torch.Tensor | None = None,
2504 + bias_vl: torch.Tensor | None = None,
2505 + image_sentinel_lo: int = 0,
2506 ) -> None:
2507 torch.ops._moe_C.topk_softplus_sqrt(
2508 topk_weights,
@@ -2513,6 +2515,8 @@ def topk_hash_softplus_sqrt(
2515 input_tokens,
2516 hash_indices_table,
2517 is_padding,
2518 + bias_vl,
2519 + image_sentinel_lo,
2520 )
2521
2522
vllm/config/model.py
+4 -1
@@ -697,7 +697,10 @@ class ModelConfig:
697 self.tokenizer_mode = "kimi_k3"
698 elif arch == "DeepseekV32ForCausalLM":
699 self.tokenizer_mode = "deepseek_v32"
700 - elif arch == "DeepseekV4ForCausalLM":
700 + elif arch in (
701 + "DeepseekV4ForCausalLM",
702 + "DeepseekV4ForConditionalGeneration",
703 + ):
704 self.tokenizer_mode = "deepseek_v4"
705 elif arch in ("InklingForCausalLM", "InklingForConditionalGeneration"):
706 self.tokenizer_mode = "inkling"
vllm/config/vllm.py
+1
@@ -77,6 +77,7 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
77 "DeepseekV32MTPModel",
78 "DeepseekV32ForCausalLM",
79 "DeepseekV4ForCausalLM",
80 + "DeepseekV4ForConditionalGeneration",
81 "DeepSeekV4MTPModel",
82 "Dots3NoteForCausalLM",
83 "Dots3NoteMTPModel",
vllm/model_executor/layers/fused_moe/layer.py
+7
@@ -131,6 +131,8 @@ def FusedMoEFactory(
131 apply_routed_scale_to_output: bool = False,
132 zero_expert_type: str | None = None,
133 hash_indices_table: torch.Tensor | None = None,
134 + bias_vl: torch.Tensor | None = None,
135 + image_sentinel_lo: int = 0,
136 runner_cls: type[MoERunner] | None = None,
137 runner_args: dict[str, Any] | None = None,
138 routed_experts_cls: type[RoutedExperts] | None = None,
@@ -199,6 +201,9 @@ def FusedMoEFactory(
201 output instead of topk_weights
202 zero_expert_type: Type of zero expert handling
203 hash_indices_table: Hash table for expert indices
204 + bias_vl: Vision routing bias for image tokens (Deepseek V4)
205 + image_sentinel_lo: First of five consecutive in-vocab image sentinel
206 + ids (0 = vision routing disabled)
207 runner_cls: Custom MoERunner class (None = use default MoERunner)
208 runner_args: Additional arguments for runner constructor
209 routed_experts_cls: Custom RoutedExperts class (None = use default)
@@ -313,6 +318,8 @@ def FusedMoEFactory(
318 zero_expert_type=zero_expert_type,
319 num_logical_experts=logical_num_experts,
320 hash_indices_table=hash_indices_table,
321 + bias_vl=bias_vl,
322 + image_sentinel_lo=image_sentinel_lo,
323 )
324
325 if params_dtype is None:
vllm/model_executor/layers/fused_moe/router/dsv4_topk.py
+30
@@ -44,8 +44,12 @@ if current_platform.is_cuda():
44 topk_weights_ptr,
45 topk_ids_ptr,
46 routed_scaling_factor,
47 + input_ids_ptr,
48 + bias_vl_ptr,
49 + image_sentinel_lo,
50 NUM_EXPERTS: tl.constexpr,
51 BLOCK_N: tl.constexpr,
52 + HAS_VL: tl.constexpr,
53 launch_pdl: tl.constexpr,
54 ):
55 row = tl.program_id(0)
@@ -54,6 +58,19 @@ if current_platform.is_cuda():
58 bias = tl.load(
59 correction_bias_ptr + expert_offsets, mask=expert_mask, other=0.0
60 ).to(tl.float32)
61 + if HAS_VL:
62 + # Image tokens carry five consecutive in-vocab sentinel ids
63 + # starting at image_sentinel_lo and use bias_vl for expert
64 + # selection instead of the regular correction bias. Ids above the
65 + # sentinel block are regular special tokens and must not match.
66 + token_id = tl.load(input_ids_ptr + row).to(tl.int64)
67 + bias_vl = tl.load(
68 + bias_vl_ptr + expert_offsets, mask=expert_mask, other=0.0
69 + ).to(tl.float32)
70 + is_image = (token_id >= image_sentinel_lo) & (
71 + token_id < image_sentinel_lo + 5
72 + )
73 + bias = tl.where(is_image, bias_vl, bias)
74
75 if launch_pdl:
76 tl.extra.cuda.gdc_wait()
@@ -101,8 +118,17 @@ def dsv4_topk(
118 correction_bias: torch.Tensor,
119 indices_dtype: torch.dtype,
120 routed_scaling_factor: float,
121 + input_ids: torch.Tensor | None = None,
122 + bias_vl: torch.Tensor | None = None,
123 + image_sentinel_lo: int = 0,
124 ) -> tuple[torch.Tensor, torch.Tensor]:
125 num_tokens, num_experts = gating_output.shape
126 + has_vl = bias_vl is not None and image_sentinel_lo > 0
127 + if bias_vl is not None:
128 + assert input_ids is not None, "bias_vl routing requires input_ids"
129 + assert bias_vl.dtype == torch.float32 and bias_vl.is_contiguous()
130 + assert bias_vl.shape == (num_experts,)
131 + assert input_ids.is_contiguous()
132 shape = (num_tokens, _TOPK)
133 topk_weights = gating_output.new_empty(shape, dtype=torch.float32)
134 topk_ids = gating_output.new_empty(shape, dtype=indices_dtype)
@@ -113,8 +139,12 @@ def dsv4_topk(
139 topk_weights,
140 topk_ids,
141 routed_scaling_factor,
142 + input_ids,
143 + bias_vl,
144 + image_sentinel_lo,
145 NUM_EXPERTS=num_experts,
146 BLOCK_N=triton.next_power_of_2(num_experts),
147 + HAS_VL=has_vl,
148 num_warps=1,
149 launch_pdl=current_platform.is_arch_support_pdl(),
150 )
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+68 -7
@@ -81,6 +81,8 @@ def vllm_topk_softplus_sqrt(
81 input_tokens: torch.Tensor | None = None,
82 hash_indices_table: torch.Tensor | None = None,
83 routed_scaling_factor: float = 1.0,
84 + bias_vl: torch.Tensor | None = None,
85 + image_sentinel_lo: int = 0,
86 ) -> tuple[torch.Tensor, ...]:
87 ops.topk_hash_softplus_sqrt(
88 topk_weights,
@@ -93,6 +95,8 @@ def vllm_topk_softplus_sqrt(
95 input_tokens,
96 hash_indices_table,
97 is_padding=_get_padding_mask(topk_indices.shape[0]),
98 + bias_vl=bias_vl,
99 + image_sentinel_lo=image_sentinel_lo,
100 )
101
102 return topk_weights, topk_indices
@@ -122,6 +126,8 @@ def fused_topk_bias(
126 input_tokens: torch.Tensor | None = None,
127 hash_indices_table: torch.Tensor | None = None,
128 routed_scaling_factor: float = 1.0,
129 + bias_vl: torch.Tensor | None = None,
130 + image_sentinel_lo: int = 0,
131 ):
132 if (
133 input_tokens is not None
@@ -130,18 +136,28 @@ def fused_topk_bias(
136 ):
137 input_tokens = input_tokens.to(dtype=hash_indices_table.dtype)
138
139 + if bias_vl is not None:
140 + # Image tokens carry five consecutive in-vocab sentinel ids starting
141 + # at image_sentinel_lo and select experts with bias_vl instead of the
142 + # regular routing path. image_sentinel_lo == 0 disables this.
143 + assert input_tokens is not None, "bias_vl routing requires input_tokens"
144 +
145 if not rocm_aiter_ops.is_fused_moe_enabled():
146 assert hidden_states.size(0) == gating_output.size(0), (
147 "Number of tokens mismatch"
148 )
149
150 output_indices_dtype = torch.int32 if indices_type is None else indices_type
139 - if scoring_func == "sqrtsoftplus" and can_use_dsv4_topk(
140 - gating_output,
141 - e_score_correction_bias,
142 - topk,
143 - renormalize,
144 - output_indices_dtype,
151 + if (
152 + scoring_func == "sqrtsoftplus"
153 + and hash_indices_table is None
154 + and can_use_dsv4_topk(
155 + gating_output,
156 + e_score_correction_bias,
157 + topk,
158 + renormalize,
159 + output_indices_dtype,
160 + )
161 ):
162 assert e_score_correction_bias is not None
163 return dsv4_topk(
@@ -149,6 +165,9 @@ def fused_topk_bias(
165 e_score_correction_bias,
166 output_indices_dtype,
167 routed_scaling_factor,
168 + input_ids=input_tokens,
169 + bias_vl=bias_vl,
170 + image_sentinel_lo=image_sentinel_lo,
171 )
172
173 M, _ = hidden_states.size()
@@ -200,6 +219,8 @@ def fused_topk_bias(
219 input_tokens,
220 hash_indices_table,
221 routed_scaling_factor,
222 + bias_vl=bias_vl,
223 + image_sentinel_lo=image_sentinel_lo,
224 )
225 else:
226 raise ValueError(f"Unsupported scoring function: {scoring_func}")
@@ -255,6 +276,8 @@ def fused_topk_bias(
276 input_tokens,
277 hash_indices_table,
278 routed_scaling_factor,
279 + bias_vl=bias_vl,
280 + image_sentinel_lo=image_sentinel_lo,
281 )
282
283 n_routed_experts = gating_output.shape[-1]
@@ -270,10 +293,39 @@ def fused_topk_bias(
293 ) + e_score_correction_bias.unsqueeze(0)
294 else:
295 scores_for_choice = scores.view(-1, n_routed_experts)
296 + image_mask = None
297 + if bias_vl is not None and image_sentinel_lo > 0:
298 + # Image tokens (five consecutive sentinel ids starting at
299 + # image_sentinel_lo) select experts with bias_vl instead of
300 + # e_score_correction_bias / the hash table. Ids above the sentinel
301 + # block are regular special tokens and must not match.
302 + assert input_tokens is not None, "bias_vl routing requires input_tokens"
303 + image_mask = (
304 + (input_tokens >= image_sentinel_lo) & (input_tokens < image_sentinel_lo + 5)
305 + ).unsqueeze(-1)
306 + text_bias = (
307 + e_score_correction_bias
308 + if e_score_correction_bias is not None
309 + else torch.zeros_like(bias_vl)
310 + )
311 + row_bias = torch.where(image_mask, bias_vl, text_bias)
312 + scores_for_choice = scores.view(-1, n_routed_experts) + row_bias
313 # For batch invariance, use sorted=True to ensure deterministic expert selection
314 if hash_indices_table is not None:
315 assert input_tokens is not None
276 - topk_indices = hash_indices_table[input_tokens]
316 + if image_mask is None:
317 + topk_indices = hash_indices_table[input_tokens]
318 + else:
319 + # Clamp sentinel rows (overwritten below) to keep the table
320 + # lookup well-defined.
321 + safe_ids = torch.where(image_mask.squeeze(-1), 0, input_tokens).to(
322 + hash_indices_table.dtype
323 + )
324 + topk_indices = hash_indices_table[safe_ids]
325 + vl_indices = torch.topk(scores_for_choice, k=topk, dim=-1)[1]
326 + topk_indices = torch.where(
327 + image_mask, vl_indices.to(topk_indices.dtype), topk_indices
328 + )
329 else:
330 use_sorted = envs.VLLM_BATCH_INVARIANT
331 topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[
@@ -306,6 +358,8 @@ class FusedTopKBiasRouter(BaseRouter):
358 hash_indices_table: torch.Tensor | None = None,
359 num_fused_shared_experts: int = 0,
360 shared_expert_weight: float = 1.0,
361 + bias_vl: torch.Tensor | None = None,
362 + image_sentinel_lo: int = 0,
363 ):
364 super().__init__(
365 top_k=top_k,
@@ -318,6 +372,11 @@ class FusedTopKBiasRouter(BaseRouter):
372 self.routed_scaling_factor = routed_scaling_factor
373 self.scoring_func = scoring_func
374 self._hash_indices_table = hash_indices_table
375 + # Vision bias: image sentinel tokens (five consecutive in-vocab ids
376 + # starting at image_sentinel_lo) select experts with bias_vl instead
377 + # of e_score_correction_bias / the hash table.
378 + self.bias_vl = bias_vl
379 + self.image_sentinel_lo = image_sentinel_lo
380 # Fused shared experts: append constant slots (ids immediately after
381 # the routed experts, [global, global+n)) routed to by every token at
382 # ``shared_expert_weight``, AFTER the routed top-k is renormalized.
@@ -357,6 +416,8 @@ class FusedTopKBiasRouter(BaseRouter):
416 input_tokens=input_ids,
417 hash_indices_table=self._hash_indices_table,
418 routed_scaling_factor=self.routed_scaling_factor,
419 + bias_vl=self.bias_vl.data if self.bias_vl is not None else None,
420 + image_sentinel_lo=self.image_sentinel_lo,
421 )
422
423 if self.num_fused_shared_experts > 0:
vllm/model_executor/layers/fused_moe/router/router_factory.py
+5
@@ -60,6 +60,9 @@ def create_fused_moe_router(
60 zero_expert_type: str | None = None,
61 num_logical_experts: int | None = None,
62 hash_indices_table: torch.Tensor | None = None,
63 + # Deepseek V4 vision routing bias parameters
64 + bias_vl: torch.Tensor | None = None,
65 + image_sentinel_lo: int = 0,
66 ) -> FusedMoERouter:
67 """
68 Factory function to create the appropriate FusedMoERouter subclass based on
@@ -211,6 +214,8 @@ def create_fused_moe_router(
214 hash_indices_table=hash_indices_table,
215 num_fused_shared_experts=num_fused_shared_experts,
216 shared_expert_weight=shared_expert_weight,
217 + bias_vl=bias_vl,
218 + image_sentinel_lo=image_sentinel_lo,
219 )
220
221 if (
vllm/model_executor/models/config.py
+1
@@ -1002,6 +1002,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
1002 "ColBERTJinaRobertaModel": JinaRobertaModelConfig,
1003 "ColQwen3_5": ColQwen3_5Config,
1004 "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig,
1005 + "DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig,
1006 "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM,
1007 "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501
1008 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501
vllm/model_executor/models/registry.py
+4
@@ -362,6 +362,10 @@ _MULTIMODAL_MODELS = {
362 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
363 "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
364 "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
365 + "DeepseekV4ForConditionalGeneration": (
366 + "vllm.models.deepseek_v4",
367 + "DeepseekV4ForConditionalGeneration",
368 + ),
369 "Dots3NoteForCausalLM": (
370 "vllm.models.dots3_note",
371 "Dots3NoteForCausalLM",
vllm/models/deepseek_v4/__init__.py
+10
@@ -20,7 +20,13 @@ if current_platform.is_rocm():
20 )
21 from .amd.model import DeepseekV4ForCausalLM
22 from .amd.mtp import DeepSeekV4MTP
23 + from .vl_stub import ( # type: ignore[assignment]
24 + DeepseekV4ForConditionalGeneration,
25 + )
26 elif current_platform.is_xpu():
27 + from .vl_stub import ( # type: ignore[assignment]
28 + DeepseekV4ForConditionalGeneration,
29 + )
30 from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment]
31 from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment]
32 from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment]
@@ -30,10 +36,14 @@ else:
36 )
37 from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment]
38 from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment]
39 + from .nvidia.vl_model import ( # type: ignore[assignment]
40 + DeepseekV4ForConditionalGeneration,
41 + )
42
43 __all__ = [
44 "DSparkDeepseekV4ForCausalLM",
45 "DeepSeekV4MTP",
46 "DeepseekV4FP8Config",
47 "DeepseekV4ForCausalLM",
48 + "DeepseekV4ForConditionalGeneration",
49 ]
vllm/models/deepseek_v4/attention.py
+8
@@ -205,6 +205,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
205 self.n_groups = config.o_groups
206 self.n_local_groups = self.n_groups // tp_size
207 self.window_size = config.sliding_window
208 + # Vision variant: image spans are visible bidirectionally, widening
209 + # prefill SWA index rows by up to max_image_tokens columns.
210 + self.max_image_tokens = (
211 + getattr(config, "vision_max_n_token", 0)
212 + if getattr(config, "vision_n_layers", 0) > 0
213 + else 0
214 + )
215 # NOTE(zyongye) Compress ratio can't be 0
216 # we do this for because MTP layer is not included
217 # in the compress ratio list
@@ -363,6 +370,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
370 _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(
371 window_size=self.window_size,
372 block_size=self.swa_cache_layer.block_size,
373 + max_image_tokens=self.max_image_tokens,
374 )
375
376 if self.compress_ratio > 1:
vllm/models/deepseek_v4/common/mm_preprocess.py new
+512
@@ -0,0 +1,512 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Multimodal preprocessing for the DeepSeek-V4 vision variants
4 +(DeepSeek-V4-Flash-Vision-Exp).
5 +
6 +The image transform and sentinel-block construction are ported from the
7 +official repository's ``image_processor.py`` so that token counts bit-match
8 +the reference. Each ``<｜deepseek_image｜>`` placeholder in the prompt expands
9 +to a variable-length block of sentinel tokens; only positions with
10 +``type == IMAGE`` receive vision embeddings, the other sentinels are looked
11 +up from learned embedding vectors in the model.
12 +
13 +Unlike the reference (which uses out-of-vocab ids ``vocab_size + type``),
14 +the sentinel block borrows five consecutive reserved tokenizer tokens
15 +(``<|place_holder_mm_span_0431|>`` .. ``_0435|>``): they are special tokens
16 +the tokenizer never emits from plain text, so the ids stay in-vocabulary and
17 +work with stock token validation, logprobs and detokenization. The ids are
18 +pure markers — every sentinel position's embedding is overwritten with
19 +vision/sentinel vectors before the decoder layers see it, exactly like the
20 +reference's out-of-vocab scheme.
21 +"""
22 +
23 +import math
24 +from collections.abc import Mapping, Sequence
25 +from typing import Any, cast
26 +
27 +import numpy as np
28 +import torch
29 +from PIL import Image, ImageOps
30 +from transformers import BatchFeature
31 +from typing_extensions import assert_never
32 +
33 +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
34 +from vllm.inputs import MultiModalDataDict
35 +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
36 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems
37 +from vllm.multimodal.processing import (
38 + BaseDummyInputsBuilder,
39 + BaseMultiModalProcessor,
40 + BaseProcessingInfo,
41 + PromptReplacement,
42 + PromptUpdate,
43 + PromptUpdateDetails,
44 +)
45 +from vllm.multimodal.processing.processor import (
46 + MultiModalPromptUpdates,
47 + MultiModalPromptUpdatesApplyResult,
48 + PlaceholderFeaturesInfo,
49 + UpdateMode,
50 + _plan_prompt_updates,
51 +)
52 +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config
53 +
54 +IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5)
55 +COMPRESS_PAD_TO = 4
56 +
57 +IMAGE_PLACEHOLDER = "<｜deepseek_image｜>"
58 +
59 +# Sentinel roles borrow five consecutive ``<|place_holder_mm_span_XXXX|>``
60 +# tokens (reserved special tokens, never emitted from plain text). The order
61 +# must match IMAGE_START..IMAGE_END above so that
62 +# ``id == IMAGE_SENTINEL_BASE_ID + type``.
63 +IMAGE_SENTINEL_BASE_ID = 129257
64 +IMAGE_SENTINEL_TOKEN_NAMES = (
65 + "<|place_holder_mm_span_0431|>", # IMAGE_START
66 + "<|place_holder_mm_span_0432|>", # IMAGE_PAD
67 + "<|place_holder_mm_span_0433|>", # IMAGE
68 + "<|place_holder_mm_span_0434|>", # IMAGE_NEW_LINE
69 + "<|place_holder_mm_span_0435|>", # IMAGE_END
70 +)
71 +
72 +
73 +def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor:
74 + """Boolean mask for image-block sentinel positions (in-vocab ids)."""
75 + return (token_ids >= IMAGE_SENTINEL_BASE_ID) & (
76 + token_ids < IMAGE_SENTINEL_BASE_ID + len(IMAGE_SENTINEL_TOKEN_NAMES)
77 + )
78 +
79 +
80 +def validate_image_sentinel_ids(tokenizer) -> None:
81 + """Check the borrowed sentinel ids against the tokenizer."""
82 + for i, name in enumerate(IMAGE_SENTINEL_TOKEN_NAMES):
83 + token_id = tokenizer.convert_tokens_to_ids(name)
84 + if token_id != IMAGE_SENTINEL_BASE_ID + i:
85 + raise ValueError(
86 + f"Image sentinel token {name!r} has id {token_id}, expected "
87 + f"{IMAGE_SENTINEL_BASE_ID + i}; the DeepSeek-V4 vision "
88 + "sentinel block requires these consecutive reserved ids."
89 + )
90 +
91 +
92 +def grid_tokens(best_height, best_width, patch_size, downsample_ratio):
93 + """Number of LLM tokens the aligner grid occupies (N-layout, including
94 + row/align padding)."""
95 + n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio)
96 + n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio)
97 + num_tokens = n_llm_h * (n_llm_w + 1) + 2
98 + if n_llm_h % 2 == 1:
99 + num_tokens += n_llm_w + 1
100 + num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2
101 + return n_llm_h, n_llm_w, num_tokens
102 +
103 +
104 +def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):
105 + r = height / width
106 + max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5
107 + max_h_float = max_w_float * r
108 + if max_w_float < 1.0:
109 + max_w = 1
110 + max_h = (max_n_token - 2) // (max_w + 1)
111 + if max_h % 2 == 1:
112 + max_h -= 1
113 + best_width = max_w * patch_size * downsample_ratio
114 + best_height = max_h * patch_size * downsample_ratio
115 + elif max_h_float < 2.0:
116 + max_h = 2
117 + max_w = ((max_n_token - 2) // max_h) - 1
118 + assert max_w > 1
119 + best_width = max_w * patch_size * downsample_ratio
120 + best_height = max_h * patch_size * downsample_ratio
121 + else:
122 + max_w = math.floor(max_w_float)
123 + max_h = math.floor(max_h_float)
124 + if max_h % 2 == 1:
125 + max_h -= 1
126 + beta = min(
127 + max_w * patch_size * downsample_ratio / width,
128 + max_h * patch_size * downsample_ratio / height,
129 + )
130 + best_width = math.floor(width * beta / patch_size) * patch_size
131 + best_height = math.floor(height * beta / patch_size) * patch_size
132 + n_llm_h, n_llm_w, num_tokens = grid_tokens(
133 + best_height, best_width, patch_size, downsample_ratio
134 + )
135 + return n_llm_h, n_llm_w, best_height, best_width, num_tokens
136 +
137 +
138 +def safe_resize(
139 + height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token
140 +):
141 + max_n_token -= COMPRESS_PAD_TO - 1
142 + n_llm_h, n_llm_w, num_tokens = grid_tokens(
143 + best_height, best_width, patch_size, downsample_ratio
144 + )
145 + budget = max_n_token
146 + while num_tokens > max_n_token:
147 + n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio(
148 + height, width, patch_size, downsample_ratio, budget
149 + )
150 + budget -= 1
151 + return n_llm_h, n_llm_w, best_height, best_width
152 +
153 +
154 +def load_image(
155 + image: Image.Image,
156 + *,
157 + patch_size: int,
158 + downsample_ratio: int,
159 + max_n_token: int,
160 + min_pixels: int,
161 + max_wh_ratio: float | None,
162 +):
163 + """Transform one PIL image into ViT patches.
164 +
165 + Same math as the reference ``load_image``, except the image is already
166 + decoded (vLLM supplies PIL images instead of a record dict).
167 + """
168 + p = patch_size
169 + image = image.convert("RGB")
170 + width, height = image.size
171 + if max_wh_ratio is not None and width > height * max_wh_ratio:
172 + width = height * max_wh_ratio
173 + if 0 < width * height < min_pixels:
174 + ratio = (min_pixels / (width * height)) ** 0.5
175 + width = int(width * ratio)
176 + height = int(height * ratio)
177 + best_width = math.ceil(width / p) * p
178 + best_height = math.ceil(height / p) * p
179 + n_llm_h, n_llm_w, best_height, best_width = safe_resize(
180 + height, width, best_height, best_width, p, downsample_ratio, max_n_token
181 + )
182 + n_vit_h, n_vit_w = best_height // p, best_width // p
183 + if max_wh_ratio is not None and image.width >= max_wh_ratio * image.height:
184 + image = image.resize((best_width, best_height))
185 + else:
186 + image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))
187 + x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
188 + x = ((x - 0.5) / 0.5).to(torch.bfloat16)
189 + patches = (
190 + x.reshape(3, n_vit_h, p, n_vit_w, p)
191 + .permute(1, 3, 0, 2, 4)
192 + .reshape(n_vit_h * n_vit_w, 3, p, p)
193 + )
194 + return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w
195 +
196 +
197 +def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int):
198 + """Builds the N-layout token types (final order) and the aligner-row order
199 + for IMAGE slots."""
200 + compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO
201 + pad_h = n_llm_h % 2
202 + rows = n_llm_h + pad_h
203 + row_len = n_llm_w + 1
204 + pad_last = rows // 2 * row_len % 2 * 2
205 + types = torch.tensor(
206 + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h
207 + + [IMAGE_PAD] * (row_len * pad_h),
208 + dtype=torch.int64,
209 + )
210 + order = torch.arange(rows * row_len).view(rows // 2, 2, row_len)
211 + order = order.transpose(1, 2).reshape(-1)
212 + image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64)
213 + image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange(
214 + n_llm_h * n_llm_w
215 + ).view(n_llm_h, n_llm_w)
216 + perm = image_idx[order]
217 + perm = perm[perm >= 0]
218 + types = torch.cat(
219 + [
220 + torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64),
221 + torch.tensor([IMAGE_START]),
222 + types[order],
223 + torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64),
224 + torch.tensor([IMAGE_END]),
225 + ]
226 + )
227 + return types, perm
228 +
229 +
230 +def build_image_block_pad_free(n_llm_h: int, n_llm_w: int):
231 + """``build_image_block`` without the position-dependent compressor pad.
232 +
233 + ``start_pos`` is chosen so that ``compress_pad == 0``; the pad is instead
234 + prepended when the block is spliced into the final prompt, where its
235 + position is known.
236 + """
237 + return build_image_block(n_llm_h, n_llm_w, COMPRESS_PAD_TO - 1)
238 +
239 +
240 +class DeepseekV4VLImageProcessor:
241 + """Per-image transform (the PIL-input equivalent of the reference
242 + ``load_image``)."""
243 +
244 + def __init__(self, config: DeepseekV4Config) -> None:
245 + super().__init__()
246 + self.patch_size = config.vision_patch_size
247 + self.downsample_ratio = config.vision_downsample_ratio
248 + self.max_n_token = config.vision_max_n_token
249 + self.min_pixels = config.vision_min_pixels
250 + self.max_wh_ratio = config.vision_max_wh_ratio
251 +
252 + def __call__(self, image: Image.Image):
253 + return load_image(
254 + image,
255 + patch_size=self.patch_size,
256 + downsample_ratio=self.downsample_ratio,
257 + max_n_token=self.max_n_token,
258 + min_pixels=self.min_pixels,
259 + max_wh_ratio=self.max_wh_ratio,
260 + )
261 +
262 +
263 +class DeepseekV4VLProcessor:
264 + """Minimal stand-in for the HF processor of DeepSeek-V4 vision models.
265 +
266 + The official repository ships image preprocessing as plain functions in
267 + ``image_processor.py`` (no ``auto_map`` processor), so this class wraps
268 + their ports directly and the model loads without ``--trust-remote-code``.
269 +
270 + ``__call__`` returns a ``BatchFeature`` with one entry per image
271 + (flattened across images):
272 +
273 + - ``patches``: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` bf16 ViT patches.
274 + - ``vit_grid``: ``(num_images, 2)`` int64 ``[n_vit_h, n_vit_w]``.
275 + - ``llm_grid``: ``(num_images, 2)`` int64 ``[n_llm_h, n_llm_w]``.
276 + - ``perm``: concatenated per-image ``(n_llm_h * n_llm_w,)`` int64 index
277 + selecting aligner outputs into the final N-layout order.
278 + - ``types``: concatenated per-image pad-free sentinel block types;
279 + ``block ids = IMAGE_SENTINEL_BASE_ID + types``.
280 + """
281 +
282 + def __init__(self, config: DeepseekV4Config) -> None:
283 + super().__init__()
284 + self.config = config
285 + self.image_processor = DeepseekV4VLImageProcessor(config)
286 +
287 + def __call__(
288 + self,
289 + text: str | None = None,
290 + images: Sequence[Image.Image] | None = None,
291 + return_tensors: str | None = None,
292 + **kwargs: Any,
293 + ) -> BatchFeature:
294 + patches_list = []
295 + vit_grid = []
296 + llm_grid = []
297 + perm_list = []
298 + types_list = []
299 + for image in images or []:
300 + patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = self.image_processor(image)
301 + types, perm = build_image_block_pad_free(n_llm_h, n_llm_w)
302 + patches_list.append(patches)
303 + vit_grid.append((n_vit_h, n_vit_w))
304 + llm_grid.append((n_llm_h, n_llm_w))
305 + perm_list.append(perm)
306 + types_list.append(types)
307 +
308 + if not patches_list:
309 + return BatchFeature({})
310 +
311 + return BatchFeature(
312 + {
313 + "patches": torch.cat(patches_list),
314 + "vit_grid": torch.tensor(vit_grid, dtype=torch.int64),
315 + "llm_grid": torch.tensor(llm_grid, dtype=torch.int64),
316 + "perm": torch.cat(perm_list),
317 + "types": torch.cat(types_list),
318 + }
319 + )
320 +
321 +
322 +class DeepseekV4VLProcessingInfo(BaseProcessingInfo):
323 + def get_hf_config(self) -> DeepseekV4Config:
324 + return self.ctx.get_hf_config(DeepseekV4Config)
325 +
326 + def get_hf_processor(self, **kwargs: object) -> DeepseekV4VLProcessor:
327 + if kwargs:
328 + raise ValueError(f"Unexpected processor kwargs: {sorted(kwargs)}")
329 + return DeepseekV4VLProcessor(self.get_hf_config())
330 +
331 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:
332 + return {"image": None}
333 +
334 + def get_mm_max_tokens_per_item(
335 + self,
336 + seq_len: int,
337 + mm_counts: Mapping[str, int],
338 + ) -> Mapping[str, int]:
339 + # ``safe_resize`` reserves COMPRESS_PAD_TO - 1 tokens of
340 + # vision_max_n_token for the compressor-alignment pad, so the full
341 + # sentinel block is bounded by vision_max_n_token; the margin is kept
342 + # in case that reservation changes.
343 + return {
344 + "image": self.get_hf_config().vision_max_n_token + COMPRESS_PAD_TO - 1,
345 + }
346 +
347 + def get_image_placeholder_token_id(self) -> int:
348 + token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER)
349 + if token_id is None:
350 + raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}")
351 + return token_id
352 +
353 + def get_image_size_with_most_features(self) -> ImageSize:
354 + hf_config = self.get_hf_config()
355 + patch_size = hf_config.vision_patch_size
356 + downsample_ratio = hf_config.vision_downsample_ratio
357 + # A square maximizes the ViT patch count (area) within the token
358 + # budget; solve the budget-derived size directly to keep the dummy
359 + # image small.
360 + budget = hf_config.vision_max_n_token - (COMPRESS_PAD_TO - 1)
361 + side = budget * patch_size * downsample_ratio
362 + _, _, best_h, best_w, _ = solve_resize_ratio(
363 + side, side, patch_size, downsample_ratio, budget
364 + )
365 + return ImageSize(width=best_w, height=best_h)
366 +
367 +
368 +class DeepseekV4VLDummyInputsBuilder(
369 + BaseDummyInputsBuilder[DeepseekV4VLProcessingInfo]
370 +):
371 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
372 + return IMAGE_PLACEHOLDER * mm_counts.get("image", 0)
373 +
374 + def get_dummy_mm_data(
375 + self,
376 + seq_len: int,
377 + mm_counts: Mapping[str, int],
378 + mm_options: Mapping[str, BaseDummyOptions],
379 + ) -> MultiModalDataDict:
380 + size = self.info.get_image_size_with_most_features()
381 + return {
382 + "image": self._get_dummy_images(
383 + width=size.width,
384 + height=size.height,
385 + num_images=mm_counts.get("image", 0),
386 + overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
387 + ),
388 + }
389 +
390 +
391 +class DeepseekV4VLMultiModalProcessor(
392 + BaseMultiModalProcessor[DeepseekV4VLProcessingInfo]
393 +):
394 + def _get_mm_fields_config(
395 + self,
396 + hf_inputs: BatchFeature,
397 + hf_processor_mm_kwargs: Mapping[str, object],
398 + ) -> Mapping[str, MultiModalFieldConfig]:
399 + vit_grid = hf_inputs.get("vit_grid")
400 + llm_grid = hf_inputs.get("llm_grid")
401 +
402 + if vit_grid is None or llm_grid is None:
403 + empty = torch.empty(0, dtype=torch.long)
404 + patch_sizes = perm_sizes = types_sizes = empty
405 + else:
406 + patch_sizes = vit_grid.prod(-1)
407 + perm_sizes = llm_grid.prod(-1)
408 + n_llm_h, n_llm_w = llm_grid[:, 0], llm_grid[:, 1]
409 + # Pad-free block length; same formula as ``grid_tokens`` given
410 + # the LLM grid.
411 + types_sizes = (
412 + n_llm_h * (n_llm_w + 1)
413 + + 2
414 + + (n_llm_h % 2) * (n_llm_w + 1)
415 + + (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2
416 + )
417 +
418 + return {
419 + "patches": MultiModalFieldConfig.flat_from_sizes("image", patch_sizes),
420 + "vit_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
421 + "llm_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
422 + "perm": MultiModalFieldConfig.flat_from_sizes("image", perm_sizes),
423 + "types": MultiModalFieldConfig.flat_from_sizes(
424 + "image", types_sizes, keep_on_cpu=True
425 + ),
426 + }
427 +
428 + def _get_prompt_updates(
429 + self,
430 + mm_items: MultiModalDataItems,
431 + hf_processor_mm_kwargs: Mapping[str, object],
432 + out_mm_kwargs: MultiModalKwargsItems,
433 + ) -> Sequence[PromptUpdate]:
434 + image_token_id = self.info.get_image_placeholder_token_id()
435 + validate_image_sentinel_ids(self.info.get_tokenizer())
436 + image_embed_id = IMAGE_SENTINEL_BASE_ID + IMAGE
437 +
438 + def get_image_replacement(item_idx: int) -> PromptUpdateDetails:
439 + types: torch.Tensor = out_mm_kwargs["image"][item_idx]["types"].data
440 + full = (IMAGE_SENTINEL_BASE_ID + types).tolist()
441 + return PromptUpdateDetails.select_token_id(full, image_embed_id)
442 +
443 + return [
444 + PromptReplacement(
445 + modality="image",
446 + target=[image_token_id],
447 + replacement=get_image_replacement,
448 + ),
449 + ]
450 +
451 + def _apply_token_matches_with_placeholders(
452 + self,
453 + token_ids: list[int],
454 + mm_prompt_updates: MultiModalPromptUpdates,
455 + ) -> tuple[
456 + list[int],
457 + MultiModalPromptUpdatesApplyResult,
458 + Mapping[str, list[PlaceholderFeaturesInfo]],
459 + ]:
460 + # Same as the base implementation, except that each image block gets
461 + # its compressor-alignment pad (``3 - start_pos % 4`` IMAGE_PAD
462 + # sentinels) prepended while splicing: the pad depends on the block's
463 + # final position in the prompt, which is unknown when the (cacheable)
464 + # prompt updates are built.
465 + matched_updates, result = _plan_prompt_updates(token_ids, mm_prompt_updates)
466 + placeholders: dict[str, list[PlaceholderFeaturesInfo]] = {
467 + modality: [] for modality in mm_prompt_updates
468 + }
469 +
470 + pad_id = IMAGE_SENTINEL_BASE_ID + IMAGE_PAD
471 +
472 + new_token_ids = list[int]()
473 + prev_end_idx = 0
474 + for matched_update in matched_updates:
475 + update = matched_update.update
476 + match = matched_update.match
477 +
478 + if update.mode == UpdateMode.INSERT:
479 + end_idx_to_insert = match.end_idx
480 + elif update.mode == UpdateMode.REPLACE:
481 + end_idx_to_insert = match.start_idx
482 + else:
483 + assert_never(update.mode)
484 +
485 + new_token_ids.extend(token_ids[prev_end_idx:end_idx_to_insert])
486 + start_idx = len(new_token_ids)
487 +
488 + tokens = list(update.content.full)
489 + if tokens and update.modality == "image":
490 + compress_pad = COMPRESS_PAD_TO - 1 - start_idx % COMPRESS_PAD_TO
491 + tokens = [pad_id] * compress_pad + tokens
492 +
493 + if tokens:
494 + content_is_embed = update.content.is_embed
495 + is_embed = (
496 + content_is_embed(tokens) if content_is_embed is not None else None
497 + )
498 + placeholders[update.modality].append(
499 + PlaceholderFeaturesInfo(
500 + modality=update.modality,
501 + item_idx=update.item_idx,
502 + start_idx=start_idx,
503 + tokens=tokens,
504 + is_embed=is_embed,
505 + )
506 + )
507 + new_token_ids.extend(tokens)
508 +
509 + prev_end_idx = match.end_idx
510 +
511 + new_token_ids.extend(token_ids[prev_end_idx:])
512 + return new_token_ids, result, placeholders
vllm/models/deepseek_v4/common/ops/cache_utils.py
+107 -27
@@ -536,10 +536,16 @@ def combine_topk_swa_indices(
536 M: int,
537 N: int,
538 out: tuple[torch.Tensor, torch.Tensor] | None = None,
539 + left_visible: torch.Tensor | None = None,
540 + right_visible: torch.Tensor | None = None,
541 + max_image_tokens: int = 0,
542 ) -> tuple[torch.Tensor, torch.Tensor]:
543 num_tokens = topk_indices.shape[0]
544 + # max_image_tokens widens the SWA column region for in-image
545 + # bidirectional visibility; the width is fixed per model so the
546 + # caller-provided workspace stays valid on image-free batches.
547 combined_topk = (
542 - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
548 + (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1)
549 // _SPARSE_PREFILL_TOPK_ALIGNMENT
550 * _SPARSE_PREFILL_TOPK_ALIGNMENT
551 )
@@ -568,6 +574,9 @@ def combine_topk_swa_indices(
574 TOP_K=topk,
575 COMPRESS_RATIO=compress_ratio,
576 WINDOW_SIZE=window_size,
577 + left_visible=left_visible,
578 + right_visible=right_visible,
579 + max_image_tokens=max_image_tokens,
580 )
581 return combined_indices, combined_lens
582
@@ -628,6 +637,7 @@ class CombineTopkSwaIndicesKernel(
637 TOP_K: int
638 COMPRESS_RATIO: int
639 WINDOW_SIZE: int
640 + IMAGE_WIDTH: int
641 PADDED_TOP_K: int
642 input_variant: TritonPointerInputVariant
643
@@ -649,11 +659,14 @@ class CombineTopkSwaIndicesKernel(
659 query_start_loc_ptr,
660 seq_lens_ptr,
661 gather_lens_ptr,
662 + left_visible_ptr,
663 + right_visible_ptr,
664 M,
665 N,
666 TOP_K: tl.constexpr,
667 COMPRESS_RATIO: tl.constexpr,
668 WINDOW_SIZE: tl.constexpr,
669 + IMAGE_WIDTH: tl.constexpr,
670 PADDED_TOP_K: tl.constexpr,
671 ):
672 batch_idx = tl.program_id(0)
@@ -682,7 +695,19 @@ class CombineTopkSwaIndicesKernel(
695 token_idx_in_query = token_idx - query_start
696 pos = start_pos + token_idx_in_query
697 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
685 - swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
698 + if IMAGE_WIDTH > 0:
699 + # In-image bidirectional visibility: the window starts up to
700 + # max(left - (window - 1), 0) positions earlier and extends
701 + # `right` positions past the query token (both 0 outside
702 + # image spans, reducing to the plain causal window).
703 + left = tl.load(left_visible_ptr + token_idx)
704 + right = tl.load(right_visible_ptr + token_idx)
705 + else:
706 + left = 0
707 + right = 0
708 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
709 + swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
710 + swa_len = pos + right - swa_start + 1
711
712 offset = tl.arange(0, PADDED_TOP_K)
713 mask = offset < topk_len
@@ -695,18 +720,19 @@ class CombineTopkSwaIndicesKernel(
720 topk_indices + M * batch_idx,
721 mask=mask,
722 )
698 - offset = tl.arange(0, WINDOW_SIZE)
723 # Index into gathered buffer: N + (position - gather_start)
700 - # For positions [pos - swa_len + 1, pos], the buffer indices are:
701 - # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start]
702 - tl.store(
703 - combined_indices_ptr
704 - + token_idx * combined_indices_stride
705 - + topk_len
706 - + offset,
707 - M * batch_idx + N + offset + pos - swa_len + 1 - gather_start,
708 - mask=offset < swa_len,
709 - )
724 + # For positions [swa_start, pos + right], the buffer indices are:
725 + # [N + swa_start - gather_start, N + pos + right - gather_start]
726 + for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE):
727 + swa_offset = i + tl.arange(0, WINDOW_SIZE)
728 + tl.store(
729 + combined_indices_ptr
730 + + token_idx * combined_indices_stride
731 + + topk_len
732 + + swa_offset,
733 + M * batch_idx + N + swa_start + swa_offset - gather_start,
734 + mask=swa_offset < swa_len,
735 + )
736
737 combined_len = topk_len + swa_len
738 tl.store(combined_lens_ptr + token_idx, combined_len)
@@ -722,6 +748,7 @@ class CombineTopkSwaIndicesKernel(
748 topk: int,
749 compress_ratio: int,
750 WINDOW_SIZE: int,
751 + image_width: int = 0,
752 ) -> CompileKey:
753 padded_topk = next_power_of_2(topk_width)
754 input_variant = TritonPointerInputVariant.from_alignment(
@@ -734,6 +761,7 @@ class CombineTopkSwaIndicesKernel(
761 TOP_K=topk,
762 COMPRESS_RATIO=compress_ratio,
763 WINDOW_SIZE=WINDOW_SIZE,
764 + IMAGE_WIDTH=image_width,
765 PADDED_TOP_K=padded_topk,
766 input_variant=input_variant,
767 )
@@ -743,10 +771,19 @@ class CombineTopkSwaIndicesKernel(
771 return []
772
773 window_size = _hf_config_int(vllm_config, "sliding_window", 128)
774 + image_width = (
775 + _hf_config_int(vllm_config, "vision_max_n_token", 0)
776 + if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0
777 + else 0
778 + )
779 + # Warm both the plain-window variant (batches without image spans)
780 + # and the in-image bidirectional variant.
781 + image_widths = [0, image_width] if image_width > 0 else [0]
782 return self._trace_dispatch(self.dispatch)(
783 _DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS,
784 _COMBINE_TOPK_SWA_POINTER_INPUTS,
785 WINDOW_SIZE=window_size,
786 + image_width=image_widths,
787 )
788
789 def compile(self, compile_key: CompileKey) -> None:
@@ -763,11 +800,14 @@ class CombineTopkSwaIndicesKernel(
800 input_variant.pointer("query_start_loc", torch.int32),
801 input_variant.pointer("seq_lens", torch.int32),
802 input_variant.pointer("gather_lens", torch.int32),
803 + int32_ptr,
804 + int32_ptr,
805 1, # do not specialize M
806 1, # do not specialize N
807 TOP_K=compile_key.TOP_K,
808 COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
809 WINDOW_SIZE=compile_key.WINDOW_SIZE,
810 + IMAGE_WIDTH=compile_key.IMAGE_WIDTH,
811 PADDED_TOP_K=compile_key.PADDED_TOP_K,
812 grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS),
813 )
@@ -786,8 +826,13 @@ class CombineTopkSwaIndicesKernel(
826 TOP_K: int,
827 COMPRESS_RATIO: int,
828 WINDOW_SIZE: int,
829 + left_visible: torch.Tensor | None = None,
830 + right_visible: torch.Tensor | None = None,
831 + max_image_tokens: int = 0,
832 ) -> None:
833 num_reqs = seq_lens.shape[0]
834 + has_image = left_visible is not None
835 + image_width = max_image_tokens if has_image else 0
836 self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)](
837 combined_indices,
838 combined_indices.stride(0),
@@ -797,11 +842,14 @@ class CombineTopkSwaIndicesKernel(
842 query_start_loc,
843 seq_lens,
844 gather_lens,
845 + left_visible if has_image else topk_indices,
846 + right_visible if has_image else topk_indices,
847 M,
848 N,
849 TOP_K=TOP_K,
850 COMPRESS_RATIO=COMPRESS_RATIO,
851 WINDOW_SIZE=WINDOW_SIZE,
852 + IMAGE_WIDTH=image_width,
853 PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]),
854 )
855
@@ -828,20 +876,33 @@ def build_flashinfer_mixed_sparse_indices(
876 decode_is_valid_token: torch.Tensor | None = None,
877 swa_block_span: int | None = None,
878 compressed_block_span: int | None = None,
879 + prefill_left_visible: torch.Tensor | None = None,
880 + prefill_right_visible: torch.Tensor | None = None,
881 + max_image_tokens: int = 0,
882 ) -> tuple[torch.Tensor, torch.Tensor]:
883 """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches.
884
834 - Produces ``sparse_indices`` of shape ``[num_tokens, swa_index_width +
835 - padded_topk]`` (the first ``swa_index_width`` columns are SWA slot ids, the
885 + Produces ``sparse_indices`` of shape ``[num_tokens, swa_total_width +
886 + padded_topk]`` (the first ``swa_total_width`` columns are SWA slot ids, the
887 rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length
888 per token). Decode tokens read precomputed SWA/compressed indices; prefill
889 tokens derive their SWA window from the position and translate local
890 compressed indices to global slots via the block tables.
891 +
892 + When ``prefill_left_visible``/``prefill_right_visible`` are given (vision
893 + variant), the SWA column region widens by ``max_image_tokens`` and prefill
894 + tokens inside an image span get a bidirectionally widened window; decode
895 + rows are padded with -1 across the extra columns.
896 """
897 assert decode_swa_indices.dtype == torch.int32
898 assert decode_swa_indices.dim() == 2
899 swa_index_width = decode_swa_indices.shape[-1]
900 assert swa_index_width >= window_size
901 + has_image = prefill_left_visible is not None
902 + image_width = max_image_tokens if has_image else 0
903 + swa_total_width = swa_index_width + image_width
904 + if has_image:
905 + assert prefill_right_visible is not None
906 if decode_compressed_topk_lens is not None:
907 assert decode_compressed_topk_lens.dtype == torch.int32
908 assert prefill_topk_indices.dtype == torch.int32
@@ -890,7 +951,7 @@ def build_flashinfer_mixed_sparse_indices(
951 padded_topk = max(topk, decode_compressed_topk)
952 padded_topk = (padded_topk + 3) // 4 * 4
953 sparse_indices = torch.empty(
893 - (num_tokens, swa_index_width + padded_topk),
954 + (num_tokens, swa_total_width + padded_topk),
955 dtype=torch.int32,
956 device=decode_swa_indices.device,
957 )
@@ -924,6 +985,8 @@ def build_flashinfer_mixed_sparse_indices(
985 decode_is_valid_token,
986 prefill_topk_indices,
987 prefill_topk_indices.stride(0),
988 + prefill_left_visible if has_image else token_to_req_indices,
989 + prefill_right_visible if has_image else token_to_req_indices,
990 query_start_loc,
991 seq_lens,
992 token_to_req_indices,
@@ -938,6 +1001,7 @@ def build_flashinfer_mixed_sparse_indices(
1001 NUM_DECODE_TOKENS=num_decode_tokens,
1002 WINDOW_SIZE=window_size,
1003 SWA_INDEX_WIDTH=swa_index_width,
1004 + IMAGE_WIDTH=image_width,
1005 COMPRESS_RATIO=compress_ratio,
1006 TOP_K=topk,
1007 PADDED_TOP_K=padded_topk,
@@ -992,6 +1056,8 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1056 decode_is_valid_token_ptr,
1057 prefill_topk_indices_ptr,
1058 prefill_topk_stride,
1059 + left_visible_ptr,
1060 + right_visible_ptr,
1061 query_start_loc_ptr,
1062 seq_lens_ptr,
1063 token_to_req_indices_ptr,
@@ -1006,6 +1072,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1072 NUM_DECODE_TOKENS,
1073 WINDOW_SIZE: tl.constexpr,
1074 SWA_INDEX_WIDTH: tl.constexpr,
1075 + IMAGE_WIDTH: tl.constexpr,
1076 COMPRESS_RATIO: tl.constexpr,
1077 TOP_K: tl.constexpr,
1078 PADDED_TOP_K: tl.constexpr,
@@ -1017,14 +1084,18 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1084 TOPK_BLOCK_SIZE: tl.constexpr,
1085 ):
1086 token_idx = tl.program_id(0)
1087 + # Total SWA column count; > SWA_INDEX_WIDTH only for in-image
1088 + # bidirectional visibility (vision variant), where the extra columns are
1089 + # -1-padded for decode rows.
1090 + SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH
1091
1092 if token_idx < NUM_DECODE_TOKENS:
1022 - for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
1093 + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1094 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1024 - mask = offset < SWA_INDEX_WIDTH
1095 + mask = offset < SWA_TOTAL_WIDTH
1096 values = tl.load(
1097 decode_swa_indices_ptr + token_idx * decode_swa_stride + offset,
1027 - mask=mask,
1098 + mask=offset < SWA_INDEX_WIDTH,
1099 other=-1,
1100 )
1101 values = _remap_flashinfer_index(values, swa_block_size, swa_block_span)
@@ -1067,7 +1138,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1138 tl.store(
1139 sparse_indices_ptr
1140 + token_idx * sparse_indices_stride
1070 - + SWA_INDEX_WIDTH
1141 + + SWA_TOTAL_WIDTH
1142 + offset,
1143 values,
1144 mask=mask,
@@ -1081,7 +1152,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1152 else:
1153 compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32)
1154
1084 - tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len)
1155 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len)
1156 return
1157
1158 prefill_idx = token_idx - NUM_DECODE_TOKENS
@@ -1093,13 +1164,22 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1164 start_pos = seq_len - query_len
1165 token_idx_in_query = token_idx - query_start
1166 pos = start_pos + token_idx_in_query
1096 - swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
1097 - swa_start_pos = pos - swa_len + 1
1167 + if IMAGE_WIDTH > 0:
1168 + # In-image bidirectional visibility: window starts up to
1169 + # max(left - (window - 1), 0) earlier and extends `right` past pos.
1170 + left = tl.load(left_visible_ptr + token_idx)
1171 + right = tl.load(right_visible_ptr + token_idx)
1172 + else:
1173 + left = 0
1174 + right = 0
1175 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0)
1176 + swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0)
1177 + swa_len = pos + right - swa_start_pos + 1
1178 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
1179
1100 - for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE):
1180 + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE):
1181 offset = i + tl.arange(0, WINDOW_BLOCK_SIZE)
1102 - mask = offset < SWA_INDEX_WIDTH
1182 + mask = offset < SWA_TOTAL_WIDTH
1183 pos_offset = swa_start_pos + offset
1184 block_indices = pos_offset // swa_block_size
1185 block_numbers = tl.load(
@@ -1143,10 +1223,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel(
1223 tl.store(
1224 sparse_indices_ptr
1225 + token_idx * sparse_indices_stride
1146 - + SWA_INDEX_WIDTH
1226 + + SWA_TOTAL_WIDTH
1227 + offset,
1228 slot_ids,
1229 mask=mask,
1230 )
1231
1152 - tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len)
1232 + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len)
vllm/models/deepseek_v4/common/vision.py new
+232
@@ -0,0 +1,232 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek-V4 vision tower (ViT + aligner) with TP-sharded linears.
4 +
5 +Ported from the official reference implementation
6 +(deepseek-ai/DeepSeek-V4-Flash-Vision-Exp). Weight names match the HF
7 +checkpoint so no renaming is needed at load time. Attention and MLP weights
8 +are tensor-parallel sharded (replicated when the vision head count is not
9 +divisible by TP size, or under ``--mm-encoder-tp-mode data``); the patch
10 +embed and norms are replicated, so the residual stream is full-width on
11 +every rank.
12 +"""
13 +
14 +from functools import lru_cache
15 +
16 +import torch
17 +import torch.nn.functional as F
18 +from torch import nn
19 +
20 +from vllm.distributed import get_tensor_model_parallel_world_size
21 +from vllm.model_executor.layers.activation import SiluAndMul
22 +from vllm.model_executor.layers.attention.mm_encoder_attention import (
23 + MMEncoderAttention,
24 +)
25 +from vllm.model_executor.layers.linear import (
26 + ColumnParallelLinear,
27 + MergedColumnParallelLinear,
28 + QKVParallelLinear,
29 + ReplicatedLinear,
30 + RowParallelLinear,
31 +)
32 +from vllm.model_executor.models.vision import is_vit_use_data_parallel
33 +
34 +
35 +@lru_cache(8)
36 +def get_vision_cos_sin(
37 + n_h: int, n_w: int, dim: int, theta: float
38 +) -> tuple[torch.Tensor, torch.Tensor]:
39 + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
40 + hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w)
41 + wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w)
42 + freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float()
43 + freqs = (freqs * inv_freq).flatten(1)
44 + return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1)
45 +
46 +
47 +def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
48 + dtype = x.dtype
49 + x1, x2 = x.float().chunk(2, dim=-1)
50 + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype)
51 +
52 +
53 +class DeepseekV4RMSNorm(nn.Module):
54 + def __init__(self, dim: int, eps: float = 1e-6):
55 + super().__init__()
56 + self.eps = eps
57 + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
58 +
59 + def forward(self, x: torch.Tensor) -> torch.Tensor:
60 + dtype = x.dtype
61 + x = x.float()
62 + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps)
63 + return (self.weight * x).to(dtype)
64 +
65 +
66 +class DeepseekV4PatchEmbed(nn.Module):
67 + def __init__(self, config):
68 + super().__init__()
69 + # Replicated: the residual stream is full-width on every rank.
70 + self.proj = ReplicatedLinear(
71 + 3 * config.vision_patch_size**2,
72 + config.vision_dim,
73 + bias=True,
74 + quant_config=None,
75 + )
76 +
77 + def forward(self, x: torch.Tensor) -> torch.Tensor:
78 + out, _ = self.proj(x.flatten(1))
79 + return out
80 +
81 +
82 +class DeepseekV4VisionAttention(nn.Module):
83 + def __init__(self, config, prefix: str = ""):
84 + super().__init__()
85 + # Convention from minimax_m3/siglip: compute the TP/DP choice locally
86 + # instead of threading a flag through every constructor.
87 + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads)
88 + self.tp_size = (
89 + 1 if use_data_parallel else get_tensor_model_parallel_world_size()
90 + )
91 + self.n_heads = config.vision_n_heads // self.tp_size
92 + self.head_dim = config.vision_dim // config.vision_n_heads
93 + self.hidden_size = config.vision_dim
94 + self.wqkv = QKVParallelLinear(
95 + config.vision_dim,
96 + self.head_dim,
97 + config.vision_n_heads,
98 + bias=True,
99 + quant_config=None,
100 + prefix=f"{prefix}.wqkv",
101 + disable_tp=use_data_parallel,
102 + )
103 + self.wo = RowParallelLinear(
104 + config.vision_dim,
105 + config.vision_dim,
106 + bias=True,
107 + quant_config=None,
108 + prefix=f"{prefix}.wo",
109 + disable_tp=use_data_parallel,
110 + )
111 + self.attn = MMEncoderAttention(
112 + num_heads=self.n_heads,
113 + head_size=self.head_dim,
114 + prefix=f"{prefix}.attn",
115 + )
116 +
117 + def forward(
118 + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
119 + ) -> torch.Tensor:
120 + n = x.size(0)
121 + qkv, _ = self.wqkv(x)
122 + q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1))
123 + q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d)
124 + k = apply_rotary(k, cos, sin).unsqueeze(0)
125 + # One image per call: a dense batch, no varlen packing metadata.
126 + o = self.attn(q, k, v.unsqueeze(0))
127 + out, _ = self.wo(o.reshape(n, -1))
128 + return out
129 +
130 +
131 +class DeepseekV4VisionMLP(nn.Module):
132 + def __init__(self, config, prefix: str = ""):
133 + super().__init__()
134 + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads)
135 + self.w1 = MergedColumnParallelLinear(
136 + config.vision_dim,
137 + [config.vision_inter_dim] * 2,
138 + bias=False,
139 + quant_config=None,
140 + prefix=f"{prefix}.w1",
141 + disable_tp=use_data_parallel,
142 + )
143 + self.w2 = RowParallelLinear(
144 + config.vision_inter_dim,
145 + config.vision_dim,
146 + bias=False,
147 + quant_config=None,
148 + prefix=f"{prefix}.w2",
149 + disable_tp=use_data_parallel,
150 + )
151 + self.act_fn = SiluAndMul()
152 +
153 + def forward(self, x: torch.Tensor) -> torch.Tensor:
154 + gate_up, _ = self.w1(x)
155 + out, _ = self.w2(self.act_fn(gate_up))
156 + return out
157 +
158 +
159 +class DeepseekV4VisionBlock(nn.Module):
160 + def __init__(self, config, prefix: str = ""):
161 + super().__init__()
162 + self.norm1 = DeepseekV4RMSNorm(config.vision_dim)
163 + self.attn = DeepseekV4VisionAttention(config, prefix=f"{prefix}.attn")
164 + self.norm2 = DeepseekV4RMSNorm(config.vision_dim)
165 + self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp")
166 +
167 + def forward(
168 + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
169 + ) -> torch.Tensor:
170 + x = x + self.attn(self.norm1(x), cos, sin)
171 + return x + self.mlp(self.norm2(x))
172 +
173 +
174 +class DeepseekV4ViT(nn.Module):
175 + """DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE."""
176 +
177 + def __init__(self, config):
178 + super().__init__()
179 + self.rope_dim = config.vision_dim // config.vision_n_heads // 2
180 + self.rope_theta = config.vision_rope_theta
181 + self.patch_embed = DeepseekV4PatchEmbed(config)
182 + self.blocks = nn.ModuleList(
183 + [
184 + DeepseekV4VisionBlock(config, prefix=f"blocks.{i}")
185 + for i in range(config.vision_n_layers)
186 + ]
187 + )
188 + self.norm = DeepseekV4RMSNorm(config.vision_dim)
189 +
190 + def forward(
191 + self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int
192 + ) -> torch.Tensor:
193 + x = self.patch_embed(patches)
194 + cos, sin = get_vision_cos_sin(n_vit_h, n_vit_w, self.rope_dim, self.rope_theta)
195 + cos = cos.to(device=x.device)
196 + sin = sin.to(device=x.device)
197 + for block in self.blocks:
198 + x = block(x, cos, sin)
199 + return self.norm(x)
200 +
201 +
202 +class DeepseekV4Aligner(nn.Module):
203 + """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector."""
204 +
205 + def __init__(self, config):
206 + super().__init__()
207 + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads)
208 + self.downsample_ratio = config.vision_downsample_ratio
209 + in_dim = config.vision_dim * self.downsample_ratio**2
210 + self.w1 = ColumnParallelLinear(
211 + in_dim,
212 + config.hidden_size,
213 + bias=True,
214 + quant_config=None,
215 + disable_tp=use_data_parallel,
216 + )
217 + self.w2 = RowParallelLinear(
218 + config.hidden_size,
219 + config.hidden_size,
220 + bias=True,
221 + quant_config=None,
222 + disable_tp=use_data_parallel,
223 + )
224 +
225 + def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
226 + r = self.downsample_ratio
227 + x = x.view(n_vit_h, n_vit_w, -1).permute(2, 0, 1)
228 + x = F.pad(x, (0, -n_vit_w % r, 0, -n_vit_h % r))
229 + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1)
230 + hidden, _ = self.w1(x)
231 + out, _ = self.w2(F.gelu(hidden))
232 + return out
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+4
@@ -435,6 +435,10 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
435 decode_is_valid_token=decode_is_valid_token,
436 swa_block_span=swa_block_span,
437 compressed_block_span=compressed_block_span,
438 + prefill_left_visible=swa_metadata.prefill_left_visible,
439 + prefill_right_visible=swa_metadata.prefill_right_visible,
440 + # getattr for tests that bypass __init__ via object.__new__.
441 + max_image_tokens=getattr(self, "max_image_tokens", 0),
442 )
443 if cache_key != "c4a":
444 swa_metadata.flashinfer_sparse_index_cache[cache_key] = (
vllm/models/deepseek_v4/nvidia/flashmla.py
+19 -2
@@ -119,7 +119,9 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
119 else:
120 assert self.topk_indices_buffer is not None
121 top_k = self.topk_indices_buffer.shape[-1]
122 - combined_topk = round_up(top_k + self.window_size, 128)
122 + combined_topk = round_up(
123 + top_k + self.window_size + self.max_image_tokens, 128
124 + )
125 current_workspace_manager().get_simultaneous(
126 ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16),
127 ((self.max_num_batched_tokens, combined_topk), torch.int32),
@@ -311,7 +313,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
313 )
314 assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0"
315 workspace_manager = current_workspace_manager()
314 - combined_topk = round_up(top_k + self.window_size, 128)
316 + combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128)
317 for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan:
318 chunk_size = chunk_end - chunk_start
319 workspace = workspace_manager.get_simultaneous(
@@ -369,6 +371,21 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
371 chunk_M,
372 chunk_N,
373 out=(combined_indices_out, combined_lens_out),
374 + left_visible=(
375 + swa_metadata.prefill_left_visible[
376 + num_decode_tokens + query_start : num_decode_tokens + query_end
377 + ]
378 + if swa_metadata.prefill_left_visible is not None
379 + else None
380 + ),
381 + right_visible=(
382 + swa_metadata.prefill_right_visible[
383 + num_decode_tokens + query_start : num_decode_tokens + query_end
384 + ]
385 + if swa_metadata.prefill_right_visible is not None
386 + else None
387 + ),
388 + max_image_tokens=self.max_image_tokens,
389 )
390 flash_mla_sparse_fwd(
391 q=q[query_start:query_end],
vllm/models/deepseek_v4/nvidia/model.py
+30 -1
@@ -92,6 +92,8 @@ from vllm.utils.math_utils import cdiv
92 from vllm.v1.attention.backends.registry import AttentionBackendEnum
93 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
94
95 +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
96 +
97 logger = init_logger(__name__)
98
99
@@ -798,6 +800,12 @@ class DeepseekV4MoE(nn.Module):
800
801 self.gate.e_score_correction_bias = None
802 self.gate.tid2eid = None
803 + self.gate.bias_vl = None
804 + # Image tokens borrow five consecutive reserved in-vocab ids starting
805 + # at IMAGE_SENTINEL_BASE_ID; 0 disables vision routing (text model).
806 + self.image_sentinel_lo = (
807 + IMAGE_SENTINEL_BASE_ID if getattr(config, "vision_n_layers", 0) > 0 else 0
808 + )
809 is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
810 self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
811 if is_hash_moe:
@@ -813,12 +821,25 @@ class DeepseekV4MoE(nn.Module):
821 ),
822 requires_grad=False,
823 )
816 - elif getattr(config, "topk_method", None) == "noaux_tc":
824 + if getattr(config, "topk_method", None) == "noaux_tc" and (
825 + not is_hash_moe or getattr(config, "vision_n_layers", 0) > 0
826 + ):
827 + # Vision checkpoints ship a gate bias on hash layers too (it is
828 + # unused for routing there; image tokens use bias_vl instead).
829 self.gate.e_score_correction_bias = nn.Parameter(
830 torch.empty(config.n_routed_experts, dtype=torch.float32),
831 requires_grad=False,
832 )
833
834 + if getattr(config, "vision_n_layers", 0) > 0:
835 + # Vision checkpoints route image sentinel tokens with bias_vl
836 + # instead of e_score_correction_bias / the hash table. Created on
837 + # every MoE layer, hash layers included.
838 + self.gate.bias_vl = nn.Parameter(
839 + torch.empty(config.n_routed_experts, dtype=torch.float32),
840 + requires_grad=False,
841 + )
842 +
843 if config.n_shared_experts is None:
844 self.shared_experts = None
845 else:
@@ -950,6 +971,8 @@ class DeepseekV4MoE(nn.Module):
971 routed_scaling_factor=self.routed_scaling_factor,
972 e_score_correction_bias=self.gate.e_score_correction_bias,
973 hash_indices_table=self.gate.tid2eid,
974 + bias_vl=getattr(self.gate, "bias_vl", None),
975 + image_sentinel_lo=self.image_sentinel_lo,
976 swiglu_limit=self.swiglu_limit,
977 router_logits_dtype=torch.float32,
978 enable_eplb=parallel_config.enable_eplb,
@@ -962,6 +985,10 @@ class DeepseekV4MoE(nn.Module):
985 ) -> torch.Tensor:
986 if self.gate.tid2eid is not None and input_ids is None:
987 raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
988 + # getattr for test doubles that fake the gate module.
989 + bias_vl = getattr(self.gate, "bias_vl", None)
990 + if bias_vl is not None and input_ids is None:
991 + raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.")
992
993 if not self.use_mega_moe:
994 return self._forward_fused_moe(hidden_states, input_ids)
@@ -981,6 +1008,8 @@ class DeepseekV4MoE(nn.Module):
1008 input_tokens=input_ids,
1009 hash_indices_table=self.gate.tid2eid,
1010 routed_scaling_factor=self.routed_scaling_factor,
1011 + bias_vl=bias_vl.data if bias_vl is not None else None,
1012 + image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0,
1013 )
1014 activation_clamp = (
1015 float(self.swiglu_limit) if self.swiglu_limit is not None else None
vllm/models/deepseek_v4/nvidia/mtp.py
+1 -1
@@ -178,7 +178,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):
178 inputs_embeds
179 ).unsqueeze(-2)
180 hidden_states, residual, post_mix, res_mix = self.mtp_block(
181 - positions=positions, x=hidden_states, input_ids=None
181 + positions=positions, x=hidden_states, input_ids=input_ids
182 )
183 hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)
184 if self.mtp_block.use_sequence_parallel:
vllm/models/deepseek_v4/nvidia/vl_model.py new
+333
@@ -0,0 +1,333 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp).
4 +
5 +Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``:
6 +
7 +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE
8 + sentinel positions; four learned vectors (``image_start`` / ``image_pad`` /
9 + ``image_newline`` / ``image_end``) fill the remaining sentinel positions.
10 +- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor
11 + in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved
12 + in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>``
13 + (see ``common/vision.py`` for the tower itself).
14 +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before
15 + its hyper-connection stream expansion. Raw ``input_ids`` still flow into
16 + the model so the MoE router can apply ``bias_vl`` to image tokens
17 + (``requires_raw_input_tokens``).
18 +"""
19 +
20 +from collections.abc import Iterable
21 +
22 +import torch
23 +from torch import nn
24 +
25 +from vllm.model_executor.models.interfaces import (
26 + MultiModalEmbeddings,
27 + SupportsEagle3,
28 + SupportsMultiModal,
29 + SupportsPP,
30 +)
31 +from vllm.model_executor.models.utils import (
32 + AutoWeightsLoader,
33 + WeightsMapper,
34 + init_vllm_registered_model,
35 + maybe_prefix,
36 +)
37 +from vllm.multimodal import MULTIMODAL_REGISTRY
38 +
39 +from ..common.mm_preprocess import (
40 + IMAGE_PLACEHOLDER,
41 + IMAGE_SENTINEL_BASE_ID,
42 + DeepseekV4VLDummyInputsBuilder,
43 + DeepseekV4VLMultiModalProcessor,
44 + DeepseekV4VLProcessingInfo,
45 + image_sentinel_mask,
46 +)
47 +from ..common.vision import DeepseekV4Aligner, DeepseekV4ViT
48 +from .model import _make_deepseek_v4_weights_mapper
49 +
50 +
51 +def _make_deepseek_v4_vl_weights_mapper(
52 + expert_dtype: str, image_enabled: bool
53 +) -> WeightsMapper:
54 + """Text-checkpoint mapping rules re-rooted under ``language_model.``."""
55 + base = _make_deepseek_v4_weights_mapper(expert_dtype)
56 + orig_to_new_prefix: dict[str, str | None] = {
57 + "layers.": "language_model.model.layers.",
58 + "embed.": "language_model.model.embed.",
59 + "norm.": "language_model.model.norm.",
60 + "hc_head": "language_model.model.hc_head",
61 + "mtp.": "language_model.model.mtp.",
62 + }
63 + if not image_enabled:
64 + orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None})
65 + return WeightsMapper(
66 + orig_to_new_prefix=orig_to_new_prefix,
67 + orig_to_new_regex=base.orig_to_new_regex,
68 + orig_to_new_suffix={
69 + "head.weight": "language_model.lm_head.weight",
70 + "embed.weight": "embed_tokens.weight",
71 + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
72 + },
73 + orig_to_new_substr={
74 + ".shared_experts.w2": ".shared_experts.down_proj",
75 + # The MTP/DSpark draft heads are not supported for the vision
76 + # variant; drop their weights.
77 + "mtp.": None,
78 + },
79 + )
80 +
81 +
82 +@MULTIMODAL_REGISTRY.register_processor(
83 + DeepseekV4VLMultiModalProcessor,
84 + info=DeepseekV4VLProcessingInfo,
85 + dummy_inputs=DeepseekV4VLDummyInputsBuilder,
86 +)
87 +class DeepseekV4ForConditionalGeneration(
88 + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3
89 +):
90 + """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower.
91 +
92 + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters)
93 + delegates through ``language_model`` via the protocol defaults.
94 + """
95 +
96 + # The MoE router needs raw token ids to detect image sentinel tokens
97 + # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl.
98 + requires_raw_input_tokens = True
99 +
100 + @classmethod
101 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:
102 + if modality == "image":
103 + return IMAGE_PLACEHOLDER
104 + raise ValueError(f"Unsupported modality: {modality!r}")
105 +
106 + def __init__(self, *, vllm_config, prefix: str = "") -> None:
107 + super().__init__()
108 + model_config = vllm_config.model_config
109 + config = model_config.hf_config
110 + self.config = config
111 + self.multimodal_config = model_config.multimodal_config
112 + assert self.multimodal_config is not None
113 +
114 + image_enabled = (
115 + config.vision_n_layers > 0
116 + and self.multimodal_config.get_limit_per_prompt("image") > 0
117 + )
118 + with self._mark_tower_model(vllm_config, {"image"}):
119 + self.vision: DeepseekV4ViT | None = None
120 + self.aligner: DeepseekV4Aligner | None = None
121 + self.image_start: nn.Parameter | None = None
122 + self.image_end: nn.Parameter | None = None
123 + self.image_newline: nn.Parameter | None = None
124 + self.image_pad: nn.Parameter | None = None
125 + if image_enabled:
126 + self.vision = DeepseekV4ViT(config)
127 + self.aligner = DeepseekV4Aligner(config)
128 + for name in (
129 + "image_start",
130 + "image_end",
131 + "image_newline",
132 + "image_pad",
133 + ):
134 + setattr(
135 + self,
136 + name,
137 + nn.Parameter(
138 + torch.empty(config.hidden_size, dtype=torch.float32)
139 + ),
140 + )
141 + self.vision.to(dtype=model_config.dtype)
142 + self.aligner.to(dtype=model_config.dtype)
143 +
144 + with self._mark_language_model(vllm_config):
145 + # The arch convertor routes any config with a vision tower to
146 + # this wrapper class; mark the copy handed to the text backbone
147 + # so it resolves to DeepseekV4ForCausalLM instead of recursing
148 + # (with_hf_config deepcopies the config, the marker survives).
149 + config._dsv4_vl_inner = True # type: ignore[attr-defined]
150 + try:
151 + self.language_model = init_vllm_registered_model(
152 + vllm_config=vllm_config,
153 + hf_config=config,
154 + prefix=maybe_prefix(prefix, "language_model"),
155 + architectures=["DeepseekV4ForCausalLM"],
156 + )
157 + finally:
158 + del config._dsv4_vl_inner # type: ignore[attr-defined]
159 + # The outer mapper (see load_weights) fully resolves HF names into
160 + # this wrapper's namespace before AutoWeightsLoader strips the
161 + # "language_model." prefix and delegates to the child's load_weights,
162 + # so the child's own mapper must be a no-op. Its suffix rules are not
163 + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would
164 + # re-fire "head.weight" -> "lm_head.weight").
165 + self.language_model.hf_to_vllm_mapper = WeightsMapper()
166 + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign]
167 + self.language_model.make_empty_intermediate_tensors
168 + )
169 +
170 + expert_dtype = getattr(config, "expert_dtype", "fp4")
171 + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper(
172 + expert_dtype, image_enabled
173 + )
174 +
175 + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None:
176 + patches = kwargs.pop("patches", None)
177 + if patches is None:
178 + return None
179 + vit_grid = kwargs.pop("vit_grid", None)
180 + llm_grid = kwargs.pop("llm_grid", None)
181 + perm = kwargs.pop("perm", None)
182 + assert vit_grid is not None and llm_grid is not None and perm is not None
183 + return {
184 + "patches": patches,
185 + "vit_grid": vit_grid,
186 + "llm_grid": llm_grid,
187 + "perm": perm,
188 + }
189 +
190 + def _encode_image(
191 + self,
192 + patches: torch.Tensor,
193 + n_vit_h: int,
194 + n_vit_w: int,
195 + perm: torch.Tensor,
196 + ) -> torch.Tensor:
197 + assert self.vision is not None and self.aligner is not None
198 + image_embeds = self.aligner(
199 + self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w
200 + )
201 + # Reorder into the N-layout block order used in the prompt.
202 + return image_embeds[perm.to(image_embeds.device)]
203 +
204 + def _process_image_input(
205 + self,
206 + patches: torch.Tensor,
207 + vit_grid: torch.Tensor,
208 + llm_grid: torch.Tensor,
209 + perm: torch.Tensor,
210 + ) -> tuple[torch.Tensor, ...]:
211 + assert self.vision is not None and self.aligner is not None
212 + patches = patches.to(self.aligner.w1.weight.dtype)
213 +
214 + embeds: list[torch.Tensor] = []
215 + vit_offset = 0
216 + llm_offset = 0
217 + for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip(
218 + vit_grid.tolist(), llm_grid.tolist(), strict=True
219 + ):
220 + n_vit = n_vit_h * n_vit_w
221 + n_llm = n_llm_h * n_llm_w
222 + embeds.append(
223 + self._encode_image(
224 + patches[vit_offset : vit_offset + n_vit],
225 + n_vit_h,
226 + n_vit_w,
227 + perm[llm_offset : llm_offset + n_llm],
228 + )
229 + )
230 + vit_offset += n_vit
231 + llm_offset += n_llm
232 + return tuple(embeds)
233 +
234 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
235 + image_input = self._parse_and_validate_image_input(**kwargs)
236 + if image_input is None or self.vision is None:
237 + return []
238 + return self._process_image_input(
239 + image_input["patches"],
240 + image_input["vit_grid"],
241 + image_input["llm_grid"],
242 + image_input["perm"],
243 + )
244 +
245 + def embed_input_ids(
246 + self,
247 + input_ids: torch.Tensor,
248 + multimodal_embeddings: MultiModalEmbeddings | None = None,
249 + *,
250 + is_multimodal: torch.Tensor | None = None,
251 + ) -> torch.Tensor:
252 + from vllm.model_executor.models.utils import _merge_multimodal_embeddings
253 +
254 + # All ids are in-vocab here: image-block sentinels are borrowed
255 + # reserved tokens (their embedding rows are always overwritten below).
256 + inputs_embeds = self.language_model.embed_input_ids(input_ids)
257 +
258 + if self.image_start is not None:
259 + # Branch-free sentinel overwrite: safe inside compiled/captured
260 + # regions (no data-dependent control flow).
261 + sentinel_mask = image_sentinel_mask(input_ids)
262 + if is_multimodal is not None:
263 + # IMAGE positions get vision embeddings via the merge below.
264 + sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device)
265 + table = torch.stack(
266 + [
267 + self.image_start,
268 + self.image_pad,
269 + self.image_pad,
270 + self.image_newline,
271 + self.image_end,
272 + ]
273 + ).to(inputs_embeds.dtype)
274 + idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4)
275 + inputs_embeds = torch.where(
276 + sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds
277 + )
278 +
279 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
280 + return inputs_embeds
281 +
282 + assert is_multimodal is not None
283 + return _merge_multimodal_embeddings(
284 + inputs_embeds=inputs_embeds,
285 + multimodal_embeddings=multimodal_embeddings,
286 + is_multimodal=is_multimodal,
287 + )
288 +
289 + def forward(
290 + self,
291 + input_ids: torch.Tensor,
292 + positions: torch.Tensor,
293 + intermediate_tensors=None,
294 + inputs_embeds: torch.Tensor | None = None,
295 + **kwargs,
296 + ) -> torch.Tensor:
297 + return self.language_model(
298 + input_ids, positions, intermediate_tensors, inputs_embeds
299 + )
300 +
301 + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
302 + return self.language_model.compute_logits(hidden_states)
303 +
304 + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor:
305 + return self.language_model.compute_logits_local(hidden_states)
306 +
307 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
308 + return self.language_model.get_expert_mapping()
309 +
310 + def get_mtp_target_hidden_states(self) -> torch.Tensor | None:
311 + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model."""
312 + return self.language_model.get_mtp_target_hidden_states()
313 +
314 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
315 + # Map HF names into this wrapper's namespace up front and sort, so
316 + # the "language_model." group reaches the child loader as one
317 + # contiguous block (AutoWeightsLoader delegates per contiguous group,
318 + # and the child's load_weights finalizes fused expert weights, which
319 + # must not run on a partially loaded model).
320 + mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0])
321 + loader = AutoWeightsLoader(self)
322 + loaded_params = loader.load_weights(mapped)
323 + # The child's load_weights already ran its post-load finalization.
324 + self._weights_finalized = True
325 + return loaded_params
326 +
327 + def process_weights_after_loading(self) -> None:
328 + # Model-level post-load hook (called by the loader after any load
329 + # format). Under DummyModelLoader the child's load_weights — and
330 + # hence its finalize step — is bypassed, so run it here instead.
331 + if getattr(self, "_weights_finalized", False):
332 + return
333 + self.language_model.process_weights_after_loading()
vllm/models/deepseek_v4/vl_stub.py new
+14
@@ -0,0 +1,14 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Stub for platforms where the DeepSeek-V4 vision variant is unsupported."""
4 +
5 +from torch import nn
6 +
7 +
8 +class DeepseekV4ForConditionalGeneration(nn.Module):
9 + def __init__(self, *, vllm_config, prefix: str = ""):
10 + super().__init__()
11 + raise NotImplementedError(
12 + "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is only "
13 + "supported on NVIDIA GPUs for now."
14 + )
vllm/tokenizers/deepseek_v4_encoding.py
+35 -3
@@ -28,6 +28,10 @@ USER_SP_TOKEN = "<｜User｜>"
28 ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
29 LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>"
30
31 +# Placeholder text inlined at each image's position; the multimodal processor
32 +# later expands it into the sentinel token block.
33 +IMAGE_PLACEHOLDER = "<｜deepseek_image｜>"
34 +
35 # Task special tokens for internal classification tasks
36 DS_TASK_SP_TOKENS = {
37 "action": "<｜action｜>",
@@ -199,6 +203,30 @@ def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
203 # Message Rendering
204 # ============================================================
205
206 +def flatten_content_blocks(content: Any) -> Any:
207 + """Flatten OpenAI-style content blocks to plain text.
208 +
209 + Image blocks are inlined as IMAGE_PLACEHOLDER at their position. The image
210 + data itself travels out of band (multi_modal_data) and is matched to
211 + placeholders by order. Plain-string content passes through unchanged.
212 + """
213 + if not isinstance(content, list):
214 + return content
215 + parts: List[str] = []
216 + for block in content:
217 + if not isinstance(block, dict):
218 + parts.append(str(block))
219 + elif block.get("type") in ("image", "image_url"):
220 + parts.append(IMAGE_PLACEHOLDER)
221 + elif block.get("type") == "text":
222 + parts.append(block.get("text", ""))
223 + else:
224 + raise ValueError(
225 + f"Unsupported content block type: {block.get('type')!r}"
226 + )
227 + return "".join(parts)
228 +
229 +
230 def render_message(
231 index: int,
232 messages: List[Dict[str, Any]],
@@ -292,6 +320,8 @@ def render_message(
320 block_type = block.get("type")
321 if block_type == "text":
322 parts.append(block.get("text", ""))
323 + elif block_type in ("image", "image_url"):
324 + parts.append(IMAGE_PLACEHOLDER)
325 elif block_type == "tool_result":
326 tool_content = block.get("content", "")
327 if isinstance(tool_content, list):
@@ -307,7 +337,7 @@ def render_message(
337 parts.append(f"[Unsupported {block_type}]")
338 prompt += "\n\n".join(parts)
339 else:
310 - prompt += content or ""
340 + prompt += flatten_content_blocks(content) or ""
341
342 elif role == "latest_reminder":
343 prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content)
@@ -447,13 +477,15 @@ def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
477 "content_blocks": [tool_block],
478 })
479 elif role == "user":
450 - text_block = {"type": "text", "text": msg.get("content", "")}
480 + # Flatten OpenAI-style content blocks (images become placeholders)
481 + text_content = flatten_content_blocks(msg.get("content", ""))
482 + text_block = {"type": "text", "text": text_content}
483 if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None:
484 merged[-1]["content_blocks"].append(text_block)
485 else:
486 new_msg = {
487 "role": "user",
456 - "content": msg.get("content", ""),
488 + "content": text_content,
489 "content_blocks": [text_block],
490 }
491 # Preserve extra fields (task, wo_eos, mask, etc.)
vllm/transformers_utils/configs/deepseek_v4.py
+29
@@ -14,10 +14,39 @@ class DeepseekV4Config(PretrainedConfig):
14 rope_scaling: dict[str, Any] | None = None,
15 rope_parameters: dict[str, Any] | None = None,
16 rope_theta: float = 10000.0,
17 + vision_n_layers: int = 0,
18 + vision_dim: int = 1024,
19 + vision_n_heads: int = 16,
20 + vision_inter_dim: int = 2816,
21 + vision_patch_size: int = 14,
22 + vision_rope_theta: float = 10000.0,
23 + vision_downsample_ratio: int = 3,
24 + vision_max_n_token: int = 384,
25 + vision_min_pixels: int = 147456,
26 + vision_max_wh_ratio: float = 8,
27 **kwargs,
28 ):
29 self.max_position_embeddings = max_position_embeddings
30 self.rope_scaling = rope_scaling
31 self.rope_theta = rope_theta
32 self.rope_parameters = rope_scaling or rope_parameters
33 + # Vision tower config; vision_n_layers == 0 means text-only.
34 + self.vision_n_layers = vision_n_layers
35 + self.vision_dim = vision_dim
36 + self.vision_n_heads = vision_n_heads
37 + self.vision_inter_dim = vision_inter_dim
38 + self.vision_patch_size = vision_patch_size
39 + self.vision_rope_theta = vision_rope_theta
40 + self.vision_downsample_ratio = vision_downsample_ratio
41 + self.vision_max_n_token = vision_max_n_token
42 + self.vision_min_pixels = vision_min_pixels
43 + self.vision_max_wh_ratio = vision_max_wh_ratio
44 + # The sparse-SWA index kernels widen the window within image spans
45 + # in-kernel, so mm-prefix ranges longer than sliding_window must be
46 + # kept (they are only consumed by those kernels).
47 + self.mm_prefix_clamp_sliding_window = vision_n_layers > 0
48 + # The visibility span covers the sentinel block [IMAGE_START,
49 + # IMAGE_END]; the mm placeholder additionally carries a leading
50 + # compressor-alignment pad (see mm_preprocess.COMPRESS_PAD_TO).
51 + self.mm_prefix_span_leading_pad_modulus = 4 if vision_n_layers > 0 else 0
52 super().__init__(**kwargs)
vllm/transformers_utils/model_arch_config_convertor.py
+35
@@ -579,6 +579,40 @@ class DeepSeekMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
579 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
580
581
582 +class DeepseekV4ModelArchConfigConvertor(ModelArchConfigConvertorBase):
583 + def __init__(
584 + self,
585 + hf_config: PretrainedConfig,
586 + hf_text_config: PretrainedConfig,
587 + revision: str | None = None,
588 + ):
589 + # DeepSeek-V4-Flash-Vision-Exp ships the same architectures/model_type
590 + # as the text-only DeepSeek-V4-Flash; route to the VL wrapper class
591 + # when the config carries a vision tower. Mutate (not just override
592 + # get_architectures) because model-class resolution reads the raw
593 + # hf_config.architectures (get_model_architecture).
594 + # Only rewrite the stock text architecture: speculative-draft configs
595 + # (DSparkDraftModel / DeepSeekV4MTPModel) keep their own classes, and
596 + # the VL wrapper marks the config copy it hands to the inner text
597 + # backbone with _dsv4_vl_inner so this rewrite does not recurse.
598 + if (
599 + getattr(hf_config, "vision_n_layers", 0) > 0
600 + and getattr(hf_config, "architectures", None) == ["DeepseekV4ForCausalLM"]
601 + and not getattr(hf_config, "_dsv4_vl_inner", False)
602 + ):
603 + hf_config.architectures = ["DeepseekV4ForConditionalGeneration"]
604 + super().__init__(hf_config, hf_text_config, revision)
605 +
606 + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool:
607 + # The vision variant needs the mm-prefix plumbing: it makes the
608 + # scheduler prefill image spans atomically (disable_chunked_mm_input)
609 + # and routes per-request image ranges to the sparse-SWA metadata
610 + # builder, which widens the sliding window bidirectionally in-kernel.
611 + if not supports_multimodal:
612 + return False
613 + return getattr(self.hf_config, "vision_n_layers", 0) > 0
614 +
615 +
616 class MimoMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):
617 def get_num_hidden_layers(self) -> int:
618 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)
@@ -775,6 +809,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = {
809 "cohere_asr": CohereAsrModelArchConfigConvertor,
810 "dbrx": DbrxModelArchConfigConvertor,
811 "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,
812 + "deepseek_v4": DeepseekV4ModelArchConfigConvertor,
813 "diffusion_gemma_text": Gemma4ModelArchConfigConvertor,
814 "ernie_mtp": ErnieMTPModelArchConfigConvertor,
815 "falcon": FalconModelArchConfigConvertor,
vllm/v1/attention/backends/mla/sparse_swa.py
+280 -74
@@ -9,6 +9,7 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
9 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
10 from vllm.model_executor.warmup.jit_warmup import (
11 WarmupIntRange,
12 + zip_inputs,
13 )
14 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
15 LaunchSpec,
@@ -167,9 +168,15 @@ class DeepseekSparseSWAMetadata:
168 decode_swa_width: int = 0
169 # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill).
170 prefill_swa_indices: torch.Tensor | None = (
170 - None # [num_prefill_tokens, 1, window_size]
171 + None # [num_prefill_tokens, 1, prefill_index_width]
172 )
173 prefill_swa_lens: torch.Tensor | None = None # [num_prefill_tokens]
174 + # In-image bidirectional visibility (vision variant): per-token counts of
175 + # extra visible tokens to the left/right inside an image span. Indexed by
176 + # absolute (decode-first) token position; only prefill rows are written.
177 + # None when the model is text-only or the batch has no image spans.
178 + prefill_left_visible: torch.Tensor | None = None
179 + prefill_right_visible: torch.Tensor | None = None
180
181 # Number of decode/prefill requests/tokens (batch is reordered: decodes first)
182 num_decodes: int = 0
@@ -416,6 +423,17 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
423 assert hasattr(hf_config, "sliding_window")
424 self.window_size = hf_config.sliding_window
425
426 + # Vision variant: image spans (up to vision_max_n_token tokens) are
427 + # visible bidirectionally, so prefill index rows widen from
428 + # window_size to window_size + max_image_tokens. Text-only models keep
429 + # max_image_tokens == 0 and take the original code paths everywhere.
430 + self.max_image_tokens = (
431 + getattr(hf_config, "vision_max_n_token", 0)
432 + if getattr(hf_config, "vision_n_layers", 0) > 0
433 + else 0
434 + )
435 + self.prefill_index_width = self.window_size + self.max_image_tokens
436 +
437 # Detect which DeepseekV4 layer types this model uses so we only build a
438 # FlashMLA tile-scheduler plan for types that will actually be called.
439 # Models without compress_ratios (pure SWA) fall back to swaonly.
@@ -447,7 +465,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
465 self.prefill_swa_indices = torch.zeros(
466 max_tokens,
467 1,
450 - self.window_size,
468 + self.prefill_index_width,
469 dtype=torch.int32,
470 device=self.device,
471 )
@@ -456,6 +474,27 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
474 dtype=torch.int32,
475 device=self.device,
476 )
477 + # In-image visibility side buffers (vision variant only). The span
478 + # CSR is uploaded per step from mm_req_doc_ranges; span capacity is
479 + # bounded by max_tokens since each span holds at least one token.
480 + if self.max_image_tokens > 0:
481 + self.left_visible = torch.zeros(
482 + max_tokens, dtype=torch.int32, device=self.device
483 + )
484 + self.right_visible = torch.zeros(
485 + max_tokens, dtype=torch.int32, device=self.device
486 + )
487 + self.span_indptr = torch.zeros(
488 + self.vllm_config.scheduler_config.max_num_seqs + 1,
489 + dtype=torch.int32,
490 + device=self.device,
491 + )
492 + self.span_starts = torch.zeros(
493 + max_tokens, dtype=torch.int32, device=self.device
494 + )
495 + self.span_ends = torch.zeros(
496 + max_tokens, dtype=torch.int32, device=self.device
497 + )
498 self.is_valid_token = torch.zeros(
499 max_tokens,
500 dtype=torch.bool,
@@ -556,6 +595,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
595 decode_swa_indices,
596 self.decode_swa_lens,
597 self.window_size,
598 + decode_swa_indices.shape[-1],
599 + self.decode_swa_lens, # unused (HAS_IMAGE=False)
600 + self.decode_swa_lens, # unused (HAS_IMAGE=False)
601 query_start_loc,
602 seq_lens,
603 token_to_req_indices,
@@ -566,16 +608,42 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
608 token_offset=0,
609 )
610
611 + # Vision variant: per-token in-image visibility for prefill tokens.
612 + # Decode tokens are always past the image spans (spans are prefilled
613 + # atomically), so the decode path above never needs them.
614 + prefill_left_visible: torch.Tensor | None = None
615 + prefill_right_visible: torch.Tensor | None = None
616 + mm_ranges = common_attn_metadata.mm_req_doc_ranges
617 + if (
618 + self.max_image_tokens > 0
619 + and num_prefill_tokens > 0
620 + and mm_ranges
621 + and any(mm_ranges.values())
622 + ):
623 + prefill_left_visible, prefill_right_visible = self._build_image_visibility(
624 + common_attn_metadata.num_reqs,
625 + mm_ranges,
626 + num_decode_tokens,
627 + num_prefill_tokens,
628 + seq_lens,
629 + query_start_loc,
630 + token_to_req_indices,
631 + )
632 +
633 # Prefill SWA indices live in paged coordinates. `token_offset` lets
634 # the kernel read is_valid_token / token_to_req_indices at absolute
635 # prefill positions while writing output starting at index 0.
636 if num_prefill_tokens > 0:
637 + has_image = prefill_left_visible is not None
638 prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]
639 prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]
640 _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
641 prefill_swa_indices,
642 prefill_swa_lens,
643 self.window_size,
644 + self.prefill_index_width,
645 + prefill_left_visible if has_image else prefill_swa_lens,
646 + prefill_right_visible if has_image else prefill_swa_lens,
647 query_start_loc,
648 seq_lens,
649 token_to_req_indices,
@@ -584,6 +652,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
652 self.block_size,
653 num_tokens=num_prefill_tokens,
654 token_offset=num_decode_tokens,
655 + has_image=has_image,
656 )
657
658 # Pre-compute DeepseekV4 prefill metadata shared across all attention layers.
@@ -623,6 +692,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
692 if num_prefill_tokens > 0
693 else None
694 ),
695 + prefill_left_visible=prefill_left_visible,
696 + prefill_right_visible=prefill_right_visible,
697 block_size=self.block_size,
698 num_decodes=num_decodes,
699 num_prefills=num_prefills,
@@ -641,6 +712,50 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
712 **deepseek_v4_fields, # type: ignore[arg-type]
713 )
714
715 + def _build_image_visibility(
716 + self,
717 + num_reqs: int,
718 + mm_ranges: dict[int, list[tuple[int, int]]],
719 + num_decode_tokens: int,
720 + num_prefill_tokens: int,
721 + seq_lens: torch.Tensor,
722 + query_start_loc: torch.Tensor,
723 + token_to_req_indices: torch.Tensor,
724 + ) -> tuple[torch.Tensor, torch.Tensor]:
725 + """Upload image spans and compute per-token in-image visibility.
726 +
727 + Returns (left, right) int32 tensors indexed by absolute (decode-first)
728 + token position; only prefill rows are meaningful.
729 + """
730 + indptr = [0] * (num_reqs + 1)
731 + starts: list[int] = []
732 + ends: list[int] = []
733 + for req_idx in range(num_reqs):
734 + for span_start, span_end in mm_ranges.get(req_idx, ()):
735 + starts.append(span_start)
736 + ends.append(span_end)
737 + indptr[req_idx + 1] = len(starts)
738 + num_spans = len(starts)
739 + assert num_spans <= self.span_starts.shape[0]
740 + self.span_indptr[: num_reqs + 1].copy_(torch.tensor(indptr, dtype=torch.int32))
741 + if num_spans > 0:
742 + self.span_starts[:num_spans].copy_(torch.tensor(starts, dtype=torch.int32))
743 + self.span_ends[:num_spans].copy_(torch.tensor(ends, dtype=torch.int32))
744 + _compute_image_visibility_kernel[(num_prefill_tokens,)](
745 + self.left_visible,
746 + self.right_visible,
747 + self.span_indptr,
748 + self.span_starts,
749 + self.span_ends,
750 + query_start_loc,
751 + seq_lens,
752 + token_to_req_indices,
753 + self.max_image_tokens,
754 + token_offset=num_decode_tokens,
755 + )
756 + end = num_decode_tokens + num_prefill_tokens
757 + return self.left_visible[:end], self.right_visible[:end]
758 +
759 def update_draft_decode_metadata(
760 self,
761 metadata: DeepseekSparseSWAMetadata,
@@ -658,6 +773,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
773 metadata.decode_swa_indices,
774 metadata.decode_swa_lens,
775 metadata.decode_swa_indices.shape[-1],
776 + metadata.decode_swa_indices.shape[-1],
777 + metadata.decode_swa_lens, # unused (HAS_IMAGE=False)
778 + metadata.decode_swa_lens, # unused (HAS_IMAGE=False)
779 metadata.query_start_loc,
780 metadata.seq_lens,
781 metadata.token_to_req_indices,
@@ -755,90 +873,153 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
873 return result
874
875
876 +@triton.jit(do_not_specialize=["token_offset"])
877 +def _compute_image_visibility_kernel(
878 + left_visible_ptr,
879 + right_visible_ptr,
880 + span_indptr_ptr,
881 + span_starts_ptr,
882 + span_ends_ptr,
883 + query_start_loc_ptr,
884 + seq_lens_ptr,
885 + token_to_req_indices_ptr,
886 + max_image_tokens,
887 + token_offset,
888 +):
889 + """Per-token in-image visible counts (port of `get_image_visible`).
890 +
891 + One program per prefill token. A token at position pos inside a span
892 + [span_start, span_end] sees min(pos - span_start, max_image_tokens - 1)
893 + extra tokens to its left and min(span_end - pos, max_image_tokens) to its
894 + right; tokens outside every span get 0/0 (plain causal window).
895 + """
896 + pid = tl.program_id(0)
897 + token_idx = pid + token_offset
898 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
899 +
900 + query_start = tl.load(query_start_loc_ptr + req_idx)
901 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)
902 + seq_len = tl.load(seq_lens_ptr + req_idx)
903 + pos = seq_len - (query_end - query_start) + token_idx - query_start
904 +
905 + span_lo = tl.load(span_indptr_ptr + req_idx)
906 + span_hi = tl.load(span_indptr_ptr + req_idx + 1)
907 + left = tl.zeros((), dtype=tl.int32)
908 + right = tl.zeros((), dtype=tl.int32)
909 + for i in range(span_lo, span_hi):
910 + span_start = tl.load(span_starts_ptr + i)
911 + span_end = tl.load(span_ends_ptr + i)
912 + in_span = (pos >= span_start) & (pos <= span_end)
913 + left = tl.where(
914 + in_span, tl.minimum(pos - span_start, max_image_tokens - 1), left
915 + )
916 + right = tl.where(in_span, tl.minimum(span_end - pos, max_image_tokens), right)
917 + tl.store(left_visible_ptr + token_idx, left)
918 + tl.store(right_visible_ptr + token_idx, right)
919 +
920 +
921 # TODO(ben): unify this kernel to reduce duplication
922
923
924 +@triton.jit(
925 + do_not_specialize=[
926 + "swa_indices_stride",
927 + "block_table_stride",
928 + "token_offset",
929 + ]
930 +)
931 +def _compute_swa_indices_and_lens_kernel(
932 + swa_indices_ptr,
933 + swa_indices_stride,
934 + swa_lens_ptr,
935 + window_size,
936 + index_width,
937 + left_visible_ptr,
938 + right_visible_ptr,
939 + query_start_loc_ptr,
940 + seq_lens_ptr,
941 + token_to_req_indices_ptr,
942 + is_valid_token_ptr,
943 + block_table_ptr,
944 + block_table_stride,
945 + block_size,
946 + token_offset,
947 + HAS_IMAGE: tl.constexpr,
948 + TRITON_BLOCK_SIZE: tl.constexpr,
949 +):
950 + pid = tl.program_id(0)
951 + token_idx = pid + token_offset
952 + is_valid = tl.load(is_valid_token_ptr + token_idx)
953 + if not is_valid:
954 + tl.store(swa_lens_ptr + pid, 0)
955 + # Clear the row so a padded token cannot gather through stale indices.
956 + for i in range(0, index_width, TRITON_BLOCK_SIZE):
957 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
958 + tl.store(
959 + swa_indices_ptr + pid * swa_indices_stride + offset,
960 + -1,
961 + mask=offset < index_width,
962 + )
963 + return
964 +
965 + req_idx = tl.load(token_to_req_indices_ptr + token_idx)
966 +
967 + query_start = tl.load(query_start_loc_ptr + req_idx)
968 + query_end = tl.load(query_start_loc_ptr + req_idx + 1)
969 + query_len = query_end - query_start
970 +
971 + seq_len = tl.load(seq_lens_ptr + req_idx)
972 + prefix_len = seq_len - query_len
973 +
974 + pos = prefix_len + token_idx - query_start
975 + if HAS_IMAGE:
976 + # In-image bidirectional visibility widens the window: the window
977 + # starts up to max(left - (window - 1), 0) positions earlier and
978 + # extends `right` positions past the query token.
979 + left = tl.load(left_visible_ptr + token_idx)
980 + right = tl.load(right_visible_ptr + token_idx)
981 + else:
982 + left = 0
983 + right = 0
984 + left_add = tl.maximum(left - (window_size - 1), 0)
985 + start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0)
986 + end_pos = pos + right + 1
987 +
988 + swa_len = end_pos - start_pos
989 + tl.store(swa_lens_ptr + pid, swa_len)
990 +
991 + for i in range(0, index_width, TRITON_BLOCK_SIZE):
992 + offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
993 +
994 + pos_offset = start_pos + offset
995 + block_indices = pos_offset // block_size
996 + block_numbers = tl.load(
997 + block_table_ptr + req_idx * block_table_stride + block_indices,
998 + mask=pos_offset < end_pos,
999 + )
1000 + block_offsets = pos_offset % block_size
1001 + slot_ids = block_numbers * block_size + block_offsets
1002 +
1003 + slot_ids = tl.where(offset < swa_len, slot_ids, -1)
1004 + tl.store(
1005 + swa_indices_ptr + pid * swa_indices_stride + offset,
1006 + slot_ids,
1007 + mask=offset < index_width,
1008 + )
1009 +
1010 +
1011 class ComputeSWAIndicesAndLensKernel(
1012 VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"]
1013 ):
1014 @dataclass(frozen=True)
1015 class CompileKey:
1016 window_size: int
1017 + index_width: int
1018 + has_image: int
1019 block_size: int
1020 triton_block_size: int
1021
770 - @staticmethod
771 - @triton.jit(
772 - do_not_specialize=[
773 - "swa_indices_stride",
774 - "block_table_stride",
775 - "token_offset",
776 - ]
777 - )
778 - def kernel(
779 - swa_indices_ptr,
780 - swa_indices_stride,
781 - swa_lens_ptr,
782 - window_size,
783 - query_start_loc_ptr,
784 - seq_lens_ptr,
785 - token_to_req_indices_ptr,
786 - is_valid_token_ptr,
787 - block_table_ptr,
788 - block_table_stride,
789 - block_size,
790 - token_offset,
791 - TRITON_BLOCK_SIZE: tl.constexpr,
792 - ):
793 - pid = tl.program_id(0)
794 - token_idx = pid + token_offset
795 - is_valid = tl.load(is_valid_token_ptr + token_idx)
796 - if not is_valid:
797 - tl.store(swa_lens_ptr + pid, 0)
798 - # Clear the row so a padded token cannot gather through stale indices.
799 - for i in range(0, window_size, TRITON_BLOCK_SIZE):
800 - offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
801 - tl.store(
802 - swa_indices_ptr + pid * swa_indices_stride + offset,
803 - -1,
804 - mask=offset < window_size,
805 - )
806 - return
807 -
808 - req_idx = tl.load(token_to_req_indices_ptr + token_idx)
809 -
810 - query_start = tl.load(query_start_loc_ptr + req_idx)
811 - query_end = tl.load(query_start_loc_ptr + req_idx + 1)
812 - query_len = query_end - query_start
813 -
814 - seq_len = tl.load(seq_lens_ptr + req_idx)
815 - prefix_len = seq_len - query_len
816 -
817 - pos = prefix_len + token_idx - query_start
818 - start_pos = tl.maximum(pos - window_size + 1, 0)
819 - end_pos = pos + 1
820 -
821 - swa_len = end_pos - start_pos
822 - tl.store(swa_lens_ptr + pid, swa_len)
823 -
824 - for i in range(0, window_size, TRITON_BLOCK_SIZE):
825 - offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
826 -
827 - pos_offset = start_pos + offset
828 - block_indices = pos_offset // block_size
829 - block_numbers = tl.load(
830 - block_table_ptr + req_idx * block_table_stride + block_indices,
831 - mask=pos_offset < end_pos,
832 - )
833 - block_offsets = pos_offset % block_size
834 - slot_ids = block_numbers * block_size + block_offsets
835 -
836 - slot_ids = tl.where(offset < swa_len, slot_ids, -1)
837 - tl.store(
838 - swa_indices_ptr + pid * swa_indices_stride + offset,
839 - slot_ids,
840 - mask=offset < window_size,
841 - )
1022 + kernel = staticmethod(_compute_swa_indices_and_lens_kernel)
1023
1024 def dispatch( # type: ignore[override]
1025 self,
@@ -854,8 +1035,24 @@ class ComputeSWAIndicesAndLensKernel(
1035 *,
1036 window_size: int,
1037 block_size: int,
1038 + max_image_tokens: int = 0,
1039 ) -> list[CompileKey]:
1040 + # Decode rows always use the plain causal window. Vision models widen
1041 + # prefill rows by max_image_tokens and additionally need the in-image
1042 + # bidirectional variant for batches that contain image spans.
1043 + width_has_image = [(window_size, 0)]
1044 + if max_image_tokens > 0:
1045 + width_has_image += [
1046 + (window_size + max_image_tokens, 0),
1047 + (window_size + max_image_tokens, 1),
1048 + ]
1049 return self._trace_dispatch(self.dispatch)(
1050 + zip_inputs(
1051 + *(
1052 + dict(index_width=index_width, has_image=has_image)
1053 + for index_width, has_image in width_has_image
1054 + )
1055 + ),
1056 window_size=window_size,
1057 block_size=block_size,
1058 )
@@ -866,6 +1063,9 @@ class ComputeSWAIndicesAndLensKernel(
1063 swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
1064 swa_lens=int32_ptr,
1065 window_size=compile_key.window_size,
1066 + index_width=compile_key.index_width,
1067 + left_visible=int32_ptr,
1068 + right_visible=int32_ptr,
1069 query_start_loc=int32_ptr,
1070 seq_lens=int32_ptr,
1071 token_to_req_indices=int32_ptr,
@@ -874,6 +1074,7 @@ class ComputeSWAIndicesAndLensKernel(
1074 block_size=compile_key.block_size,
1075 num_tokens=1,
1076 token_offset=0,
1077 + has_image=bool(compile_key.has_image),
1078 )
1079
1080 @kernel_launcher
@@ -882,6 +1083,9 @@ class ComputeSWAIndicesAndLensKernel(
1083 swa_indices: torch.Tensor,
1084 swa_lens: torch.Tensor,
1085 window_size: int,
1086 + index_width: int,
1087 + left_visible: torch.Tensor,
1088 + right_visible: torch.Tensor,
1089 query_start_loc: torch.Tensor,
1090 seq_lens: torch.Tensor,
1091 token_to_req_indices: torch.Tensor,
@@ -891,10 +1095,12 @@ class ComputeSWAIndicesAndLensKernel(
1095 *,
1096 num_tokens: int,
1097 token_offset: int,
1098 + has_image: bool = False,
1099 ) -> LaunchSpec:
1100 return (num_tokens,), dict(
1101 swa_indices_stride=swa_indices.stride(0),
1102 block_table_stride=block_table.stride(0),
1103 + HAS_IMAGE=has_image,
1104 TRITON_BLOCK_SIZE=1024,
1105 )
1106
vllm/v1/worker/gpu_model_runner.py
+22 -1
@@ -2429,6 +2429,16 @@ class GPUModelRunner(
2429 _bidi_sw = getattr(hf_text_config, "sliding_window", None)
2430 _clamps_in_kernel = getattr(
2431 self.model, "mm_prefix_clamp_sliding_window", False
2432 + ) or getattr(hf_text_config, "mm_prefix_clamp_sliding_window", False)
2433 + # Some models (DeepSeek-V4 vision) define the bidirectional span
2434 + # over the whole sentinel block ([IMAGE_START, IMAGE_END]) rather
2435 + # than the embed tokens, and prepend a position-dependent
2436 + # alignment pad before the first sentinel. For those, derive the
2437 + # span from the full placeholder range and strip the pad.
2438 + # TODO(Isotr0py): Refactor mm_prefix_lm implementation
2439 + # for better readability and maintainability.
2440 + _span_pad_modulus = getattr(
2441 + hf_text_config, "mm_prefix_span_leading_pad_modulus", 0
2442 )
2443 for req_id in self.input_batch.req_ids:
2444 image_doc_ranges = []
@@ -2437,7 +2447,18 @@ class GPUModelRunner(
2447 if mm_feature.modality == "audio":
2448 continue
2449 pos_info = mm_feature.mm_position
2440 - img_doc_range = pos_info.extract_embeds_range()
2450 + if _span_pad_modulus:
2451 + pad = (
2452 + _span_pad_modulus - 1 - pos_info.offset % _span_pad_modulus
2453 + )
2454 + img_doc_range = [
2455 + (
2456 + pos_info.offset + pad,
2457 + pos_info.offset + pos_info.length - 1,
2458 + )
2459 + ]
2460 + else:
2461 + img_doc_range = pos_info.extract_embeds_range()
2462 for r in img_doc_range:
2463 if (
2464 not _clamps_in_kernel