diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 7c3ca10371..03057c3b14 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -28,7 +28,9 @@ void topk_softplus_sqrt( const std::optional& correction_bias, const std::optional& input_ids, const std::optional& tid2eid, - const std::optional& is_padding); + const std::optional& is_padding, + const std::optional& bias_vl, + int64_t image_sentinel_lo); void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output, std::optional topk_ids, diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index 976eebac3d..fbb07216d2 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -81,19 +81,92 @@ __launch_bounds__(128) __global__ int num_experts, float routed_scaling_factor, const HashIndType* input_ids, const HashIndType* tid2eid, - const bool* is_padding) { + const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int lane = threadIdx.x % 32; if (warp >= num_rows) return; const int64_t token_id = load_index_as_int64(input_ids, warp); const bool is_pad_row = is_padding != nullptr && is_padding[warp]; + // Image tokens carry five consecutive in-vocab sentinel ids starting at + // image_sentinel_lo (0 = disabled); they skip the tid2eid lookup and + // select experts by score + bias_vl instead. Ids above the sentinel block + // are regular special tokens and must not match. + const bool is_image = bias_vl != nullptr && image_sentinel_lo > 0 && + token_id >= image_sentinel_lo && + token_id < image_sentinel_lo + 5; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaGridDependencySynchronize(); #endif int expert = 0; float weight = 0.f; - if (lane < 6 && !is_pad_row) { + if (is_image && !is_pad_row) { + // Top-6 by (sqrtsoftplus(x) + bias_vl) over the full row; the routing + // weights are the raw sqrtsoftplus scores of the selected experts. + constexpr int MAX_PER_LANE = 12; // supports num_experts <= 384 + const int per_lane = num_experts / 32; + float vals[MAX_PER_LANE]; + float sels[MAX_PER_LANE]; + #pragma unroll + for (int i = 0; i < MAX_PER_LANE; ++i) { + if (i < per_lane) { + const int e = lane + 32 * i; + const float x = input[warp * num_experts + e]; + float v = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x)))); + if (isnan(v)) { + v = 0.f; + } + vals[i] = v; + sels[i] = v + bias_vl[e]; + } else { + vals[i] = 0.f; + sels[i] = -INFINITY; + } + } + int slot_expert[6]; + float slot_weight[6]; + #pragma unroll + for (int j = 0; j < 6; ++j) { + float best = -INFINITY; + int best_e = num_experts; + #pragma unroll + for (int i = 0; i < MAX_PER_LANE; ++i) { + if (sels[i] > best) { + best = sels[i]; + best_e = lane + 32 * i; + } + } + // Warp argmax reduce; ties go to the lower expert id. + #pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + const float other_best = VLLM_SHFL_XOR_SYNC(best, mask); + const int other_e = VLLM_SHFL_XOR_SYNC(best_e, mask); + if (other_best > best || (other_best == best && other_e < best_e)) { + best = other_best; + best_e = other_e; + } + } + float w = 0.f; + #pragma unroll + for (int i = 0; i < MAX_PER_LANE; ++i) { + if (lane + 32 * i == best_e) { + w = vals[i]; + sels[i] = -INFINITY; + } + } + #pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + w += VLLM_SHFL_XOR_SYNC(w, mask); + } + slot_expert[j] = best_e; + slot_weight[j] = w; + } + if (lane < 6) { + expert = slot_expert[lane]; + weight = slot_weight[lane]; + } + } else if (lane < 6 && !is_pad_row) { // only load and calculate for 6 experts expert = static_cast(tid2eid[token_id * 6 + lane]); const float x = input[warp * num_experts + expert]; @@ -127,7 +200,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, double routed_scaling_factor, const HashIndType* input_ids, const HashIndType* tid2eid, cudaStream_t stream, - const bool* is_padding) { + const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { if (num_rows == 0) return; auto* kernel = &dsv4HashTopkSoftplusSqrt; cudaLaunchConfig_t config = {}; @@ -141,7 +215,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, config.numAttrs = 1; const float scale = static_cast(routed_scaling_factor); cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows, - num_experts, scale, input_ids, tid2eid, is_padding); + num_experts, scale, input_ids, tid2eid, is_padding, + bias_vl, image_sentinel_lo); } #endif @@ -174,7 +249,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const HashIndType* input_ids, const HashIndType* tid2eid, - const bool* is_padding) { + const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { static_assert(std::is_same_v || std::is_same_v || std::is_same_v, @@ -241,6 +317,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ const bool row_is_active = finished ? !finished[thread_row] : true; const bool is_pad_row = is_padding != nullptr && is_padding[thread_row]; + // Image tokens carry five consecutive in-vocab sentinel ids starting at + // image_sentinel_lo (0 = disabled); they select experts with bias_vl + // instead of correction_bias / the tid2eid hash table. Ids above the + // sentinel block are regular special tokens and must not match. + int64_t token_id = 0; + if (input_ids != nullptr) { + token_id = load_index_as_int64(input_ids, thread_row); + } + const bool use_vl_bias = bias_vl != nullptr && image_sentinel_lo > 0 && + token_id >= image_sentinel_lo && + token_id < image_sentinel_lo + 5; + const float* row_bias = use_vl_bias ? bias_vl : correction_bias; + // We finally start setting up the read pointers for each thread. First, each // thread jumps to the start of the row it will read. const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW; @@ -323,9 +412,10 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ constexpr float threshold = 20.0f; constexpr float beta = 1.0f; - // Hash MoE path: indices are predetermined from lookup table - if constexpr (USE_HASH) { - const int64_t token_id = load_index_as_int64(input_ids, thread_row); + // Hash MoE path: indices are predetermined from lookup table. Image rows + // (sentinel ids) fall through to the score-based top-k path with bias_vl. + const bool hash_row = USE_HASH && !use_vl_bias; + if (hash_row) { const int64_t token_expert_offset = token_id * static_cast(k); if (!is_pad_row) { #pragma unroll @@ -417,13 +507,13 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ if (isnan(val)) { val = 0.f; } - if (correction_bias) { + if (row_bias) { const int group_id = ii / ELTS_PER_LDG; const int local_id = ii % ELTS_PER_LDG; const int expert_idx = first_elt_read_by_thread + group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id; - val = val + correction_bias[expert_idx]; + val = val + row_bias[expert_idx]; } row_chunk[ii] = val; } @@ -488,8 +578,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // to global memory. (This will be a single) thread per row of the // input/output matrices. const int idx = k * thread_row + k_idx; - if (correction_bias != nullptr && should_process_row) { - max_val -= correction_bias[expert]; + if (row_bias != nullptr && should_process_row) { + max_val -= row_bias[expert]; } output[idx] = max_val; indices[idx] = @@ -576,7 +666,8 @@ void topkGatingSoftplusSqrtLauncherHelper( const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, - const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) { + const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding, + const float* bias_vl, int64_t image_sentinel_lo) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = @@ -605,12 +696,13 @@ void topkGatingSoftplusSqrtLauncherHelper( cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, is_padding); + input_ids, tid2eid, is_padding, bias_vl, + image_sentinel_lo); #else kernel<<>>( input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, is_padding); + input_ids, tid2eid, is_padding, bias_vl, image_sentinel_lo); #endif }) } @@ -624,7 +716,7 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \ - stream, is_padding); + stream, is_padding, bias_vl, image_sentinel_lo); #else #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -633,14 +725,14 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream, is_padding); \ + tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \ } else if (WARP_SIZE == 32) { \ topkGatingSoftplusSqrtLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream, is_padding); \ + tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -654,14 +746,16 @@ void topkGatingSoftplusSqrtKernelLauncher( const int topk, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, const HashIndType* tid2eid, - cudaStream_t stream, const bool* is_padding) { + cudaStream_t stream, const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { #ifndef USE_ROCM if constexpr (std::is_same_v) { if (use_hash && topk == 6 && renormalize && (num_experts == 256 || num_experts == 384)) { launchDsv4HashTopk( gating_output, topk_weights, topk_indices, num_tokens, num_experts, - routed_scaling_factor, input_ids, tid2eid, stream, is_padding); + routed_scaling_factor, input_ids, tid2eid, stream, is_padding, + bias_vl, image_sentinel_lo); return; } } @@ -761,12 +855,30 @@ void dispatch_topk_softplus_sqrt_launch( const std::optional& correction_bias, const std::optional& input_ids, const std::optional& tid2eid, cudaStream_t stream, - const std::optional& is_padding) { + const std::optional& is_padding, + const std::optional& bias_vl, + int64_t image_sentinel_lo) { const float* bias_ptr = nullptr; if (correction_bias.has_value()) { bias_ptr = correction_bias.value().const_data_ptr(); } + const float* bias_vl_ptr = nullptr; + if (bias_vl.has_value()) { + const torch::stable::Tensor& bias_vl_tensor = bias_vl.value(); + STD_TORCH_CHECK(input_ids.has_value(), + "input_ids is required when bias_vl is set"); + STD_TORCH_CHECK( + bias_vl_tensor.scalar_type() == torch::headeronly::ScalarType::Float, + "bias_vl tensor must be float32"); + STD_TORCH_CHECK(bias_vl_tensor.dim() == 1, "bias_vl tensor must be 1D"); + STD_TORCH_CHECK(bias_vl_tensor.size(0) == num_experts, + "bias_vl size mismatch, expected: ", num_experts); + STD_TORCH_CHECK(bias_vl_tensor.is_contiguous(), + "bias_vl tensor must be contiguous"); + bias_vl_ptr = bias_vl_tensor.const_data_ptr(); + } + auto launch = [&](auto* topk_indices_ptr) { using OutIndType = typename std::remove_pointer::type; @@ -800,7 +912,8 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream, is_padding_ptr); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); } else { STD_TORCH_CHECK(tid2eid.value().scalar_type() == torch::headeronly::ScalarType::Int); @@ -810,7 +923,35 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream, is_padding_ptr); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); + } + } else if (bias_vl_ptr != nullptr) { + // Non-hash rows with bias_vl: image tokens (five sentinel ids starting + // at image_sentinel_lo) select experts by score + bias_vl. input_ids + // drives the per-row branch, so + // dispatch the hash-index type on its dtype. + if (input_ids.value().scalar_type() == + torch::headeronly::ScalarType::Long) { + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, false, input_ids.value().const_data_ptr(), + static_cast(nullptr), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); + } else { + STD_TORCH_CHECK(input_ids.value().scalar_type() == + torch::headeronly::ScalarType::Int); + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, false, input_ids.value().const_data_ptr(), + static_cast(nullptr), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); } } else { vllm::moe::topkGatingSoftplusSqrtKernelLauncher( @@ -818,7 +959,8 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, false, static_cast(nullptr), - static_cast(nullptr), stream, is_padding_ptr); + static_cast(nullptr), stream, is_padding_ptr, + nullptr, 0); } }; @@ -843,7 +985,9 @@ void topk_softplus_sqrt( const std::optional& correction_bias, const std::optional& input_ids, const std::optional& tid2eid, - const std::optional& is_padding) { + const std::optional& is_padding, + const std::optional& bias_vl, + int64_t image_sentinel_lo) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); @@ -857,21 +1001,21 @@ void topk_softplus_sqrt( gating_output.const_data_ptr(), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, input_ids, tid2eid, stream, - is_padding); + is_padding, bias_vl, image_sentinel_lo); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_softplus_sqrt_launch<__half>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream, is_padding); + input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream, is_padding); + input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index a321043848..800284d70f 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -20,7 +20,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, float " "routed_scaling_factor, Tensor? " - "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()"); + "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding, " + "Tensor? bias_vl=None, int image_sentinel_lo=0) -> ()"); // Calculate the result of moe by summing up the partial results // from all selected experts. topk_ids/expert_map are optional and, when diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 609ca54564..03c0b8e307 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -528,6 +528,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I+ | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ | | `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I+ | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ | | `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I+ | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ | +| `DeepseekV4ForConditionalGeneration` | DeepSeek-V4 (vision) | T + I+ | `deepseek-ai/DeepSeek-V4-Flash-Vision-Exp` | | ✅︎ | | `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + IE+ | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ | | `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I+/ V+ | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ | | `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + IE+ | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ | diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index fb2db3f2c7..b68820c0d5 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -465,3 +465,34 @@ def test_draft_model_arch_config( _assert_model_config_methods( model_config, expected, check_head_size=check_head_size ) + + +def test_deepseek_v4_convertor_splits_vision_architecture(): + """The vision checkpoint shares model_type/architectures with the + text-only DeepSeek-V4; the convertor routes it to the VL wrapper class + by rewriting hf_config.architectures (model-class resolution reads the + raw attribute).""" + from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config + from vllm.transformers_utils.model_arch_config_convertor import ( + DeepseekV4ModelArchConfigConvertor, + ) + + text_cfg = DeepseekV4Config(architectures=["DeepseekV4ForCausalLM"]) + conv = DeepseekV4ModelArchConfigConvertor(text_cfg, text_cfg) + assert conv.get_architectures() == ["DeepseekV4ForCausalLM"] + assert text_cfg.architectures == ["DeepseekV4ForCausalLM"] + + vl_cfg = DeepseekV4Config( + architectures=["DeepseekV4ForCausalLM"], vision_n_layers=32 + ) + conv = DeepseekV4ModelArchConfigConvertor(vl_cfg, vl_cfg) + assert conv.get_architectures() == ["DeepseekV4ForConditionalGeneration"] + assert vl_cfg.architectures == ["DeepseekV4ForConditionalGeneration"] + + # Speculative-draft configs (set by SpeculativeConfig.hf_config_override) + # keep their own architecture even with a vision tower in the config. + for draft_arch in ("DSparkDraftModel", "DeepSeekV4MTPModel"): + draft_cfg = DeepseekV4Config(architectures=[draft_arch], vision_n_layers=32) + conv = DeepseekV4ModelArchConfigConvertor(draft_cfg, draft_cfg) + assert conv.get_architectures() == [draft_arch] + assert draft_cfg.architectures == [draft_arch] diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index 47202f482c..ab9072576d 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -26,18 +26,46 @@ def _torch_topk_softplus_sqrt( e_score_correction_bias: torch.Tensor | None = None, input_ids: torch.Tensor | None = None, hash_indices_table: torch.Tensor | None = None, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ): scores = F.softplus(gating_output.float()).sqrt() original_scores = scores - if e_score_correction_bias is not None: - scores_for_choice = scores + e_score_correction_bias.unsqueeze(0) - else: - scores_for_choice = scores if hash_indices_table is not None: assert input_ids is not None - topk_ids = hash_indices_table[input_ids.long()] + if bias_vl is None or image_sentinel_lo == 0: + topk_ids = hash_indices_table[input_ids.long()] + else: + image_mask = (input_ids.long() >= image_sentinel_lo) & ( + input_ids.long() < image_sentinel_lo + 5 + ) + topk_ids = hash_indices_table[input_ids.long()] + vl_ids = (scores + bias_vl).topk(topk, dim=-1).indices + topk_ids = torch.where( + image_mask.unsqueeze(-1), vl_ids.to(topk_ids.dtype), topk_ids + ) else: + if bias_vl is not None and image_sentinel_lo > 0: + # Image tokens carry five consecutive in-vocab sentinel ids + # starting at image_sentinel_lo and select experts with bias_vl + # instead of the regular bias. + assert input_ids is not None + assert e_score_correction_bias is not None + image_mask = ( + (input_ids.long() >= image_sentinel_lo) + & (input_ids.long() < image_sentinel_lo + 5) + ).unsqueeze(-1) + row_bias = torch.where( + image_mask, + bias_vl.unsqueeze(0), + e_score_correction_bias.unsqueeze(0), + ) + scores_for_choice = scores + row_bias + elif e_score_correction_bias is not None: + scores_for_choice = scores + e_score_correction_bias.unsqueeze(0) + else: + scores_for_choice = scores # Match the fused kernel's deterministic tie-break: lower expert ids # win when scores are equal. torch.topk does not guarantee which tied # index it selects at the k-th boundary. @@ -368,3 +396,227 @@ def test_fused_topk_softplus_sqrt_padding( sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref) sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops) torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) + + +def _make_mixed_input_ids( + num_tokens: int, image_sentinel_lo: int, dtype: torch.dtype = torch.long +) -> torch.Tensor: + """Mix of text, image-sentinel, and above-sentinel ids. + + Every third token is an image sentinel id (cycling through the five slots + [image_sentinel_lo, image_sentinel_lo + 5)); every sixth token carries an + id just above the sentinel block, simulating the named special tokens + that follow it — those must route as regular text tokens. + """ + input_ids = torch.randint( + 0, image_sentinel_lo, (num_tokens,), dtype=torch.long, device="cuda" + ) + pos = torch.arange(num_tokens, device="cuda") + image_rows = pos % 3 == 0 + input_ids[image_rows] = image_sentinel_lo + (pos[image_rows] // 3) % 5 + above_rows = pos % 6 == 3 + input_ids[above_rows] = image_sentinel_lo + 5 + return input_ids.to(dtype) + + +def _assert_topk_matches( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights_ref: torch.Tensor, + topk_ids_ref: torch.Tensor, +) -> None: + # Different kernels may return the topk experts in different orders when + # scores tie; sort by expert id before comparing. + sorted_ref_ids, idx_ref = topk_ids_ref.sort(dim=-1) + sorted_ids, idx_ops = topk_ids.sort(dim=-1) + torch.testing.assert_close( + sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0 + ) + + sorted_w_ref = topk_weights_ref.gather(1, idx_ref) + sorted_w = topk_weights.gather(1, idx_ops.to(torch.long)) + torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="This test is skipped on non-CUDA platform.", +) +@pytest.mark.parametrize( + ("num_experts", "topk", "renormalize", "dtype"), + [ + (256, 6, True, torch.float32), # Triton dsv4 fast path + (256, 8, True, torch.bfloat16), # CUDA topk kernel + (384, 6, False, torch.half), + (512, 16, True, torch.float32), + ], +) +def test_fused_topk_softplus_sqrt_bias_vl( + num_experts: int, + topk: int, + renormalize: bool, + dtype: torch.dtype, +): + """Image sentinel tokens must select experts with bias_vl. + + Rows with ids just above the sentinel block (image_sentinel_lo + 5) + simulate the named special tokens that follow it and must route through + the regular bias path; the reference comparison covers them. + """ + torch.manual_seed(0) + num_tokens = 64 + hidden_size = 1024 + vocab_size = 128 + image_sentinel_lo = vocab_size - 8 + hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") + e_score_correction_bias = torch.randn( + (num_experts,), dtype=torch.float32, device="cuda" + ) + bias_vl = torch.randn((num_experts,), dtype=torch.float32, device="cuda") + input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo) + assert (input_ids == image_sentinel_lo + 5).any() + + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( + gating_output=gating_output, + topk=topk, + renormalize=renormalize, + routed_scaling_factor=1.5, + e_score_correction_bias=e_score_correction_bias, + input_ids=input_ids, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, + ) + + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=gating_output, + scoring_func="sqrtsoftplus", + e_score_correction_bias=e_score_correction_bias, + topk=topk, + renormalize=renormalize, + routed_scaling_factor=1.5, + input_tokens=input_ids, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, + ) + + _assert_topk_matches(topk_weights, topk_ids, topk_weights_ref, topk_ids_ref) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="This test is skipped on non-CUDA platform.", +) +@pytest.mark.parametrize( + ("num_experts", "topk", "renormalize", "dtype", "indices_dtype"), + [ + (256, 6, True, torch.float32, torch.long), # specialized hash kernel + (384, 6, True, torch.float32, torch.int32), # specialized hash kernel + (384, 8, False, torch.bfloat16, torch.long), # generic hash kernel + (512, 6, True, torch.float32, torch.long), # generic hash kernel + ], +) +def test_fused_topk_softplus_sqrt_hash_bias_vl( + num_experts: int, + topk: int, + renormalize: bool, + dtype: torch.dtype, + indices_dtype: torch.dtype, +): + """Hash MoE: text rows use tid2eid, image rows use topk(score + bias_vl). + + Sentinel ids are in-vocab; rows with ids above the sentinel block + (image_sentinel_lo + 5) must still come straight from the hash table. + """ + torch.manual_seed(0) + num_tokens = 64 + hidden_size = 1024 + vocab_size = 64 + image_sentinel_lo = vocab_size - 8 + hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") + hash_indices_table = torch.stack( + [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] + ).to(device="cuda", dtype=indices_dtype) + bias_vl = torch.randn((num_experts,), dtype=torch.float32, device="cuda") + input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo) + image_rows = (input_ids >= image_sentinel_lo) & (input_ids < image_sentinel_lo + 5) + assert image_rows.any() and (~image_rows).any() + assert (input_ids == image_sentinel_lo + 5).any() + + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( + gating_output=gating_output, + topk=topk, + renormalize=renormalize, + routed_scaling_factor=2.5, + input_ids=input_ids, + hash_indices_table=hash_indices_table, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, + ) + + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=gating_output, + scoring_func="sqrtsoftplus", + e_score_correction_bias=None, + topk=topk, + renormalize=renormalize, + input_tokens=input_ids, + hash_indices_table=hash_indices_table, + routed_scaling_factor=2.5, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, + ) + + _assert_topk_matches(topk_weights, topk_ids, topk_weights_ref, topk_ids_ref) + + # Text rows must come straight from the hash table, image rows from + # topk(score + bias_vl). + text_ids = hash_indices_table[input_ids[~image_rows].long()] + torch.testing.assert_close( + topk_ids[~image_rows].to(text_ids.dtype), text_ids, atol=0, rtol=0 + ) + + +@pytest.mark.skipif( + not current_platform.is_cuda(), + reason="The DeepSeek V4 fast path is CUDA-only.", +) +def test_dsv4_fast_topk_bias_vl(): + torch.manual_seed(0) + num_tokens = 33 + num_experts = 256 + vocab_size = 128 + image_sentinel_lo = vocab_size - 8 + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") + bias_vl = torch.randn(num_experts, dtype=torch.float32, device="cuda") + input_ids = _make_mixed_input_ids(num_tokens, image_sentinel_lo) + + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( + gating_output=gating_output, + topk=6, + renormalize=True, + routed_scaling_factor=1.5, + e_score_correction_bias=correction_bias, + input_ids=input_ids, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, + ) + topk_weights, topk_ids = dsv4_topk( + gating_output, + correction_bias, + torch.int64, + 1.5, + input_ids=input_ids, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, + ) + + assert topk_ids.dtype == torch.int64 + torch.testing.assert_close(topk_ids_ref.to(torch.int64), topk_ids, atol=0, rtol=0) + torch.testing.assert_close(topk_weights_ref, topk_weights, atol=2e-5, rtol=2e-5) diff --git a/tests/models/multimodal/processing/test_tensor_schema.py b/tests/models/multimodal/processing/test_tensor_schema.py index 87cf65b8f9..8dc59fd4f5 100644 --- a/tests/models/multimodal/processing/test_tensor_schema.py +++ b/tests/models/multimodal/processing/test_tensor_schema.py @@ -164,6 +164,11 @@ def test_model_tensor_schema(model_id: str): "Kimi-K2.5's offline inference has issues about vision chunks. Fix later." ) + if model_id == "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp" and not ( + current_platform.is_cuda() + ): + pytest.skip("Deepseek V4 is only supported on CUDA") + model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id) model_info.check_available_online(on_fail="skip") model_info.check_transformers_version( diff --git a/tests/models/registry.py b/tests/models/registry.py index 9d4de9b761..e578edf59a 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -833,6 +833,9 @@ _MULTIMODAL_EXAMPLE_MODELS = { "DeepseekOCR2ForCausalLM": _HfExamplesInfo( "deepseek-ai/DeepSeek-OCR-2", ), + "DeepseekV4ForConditionalGeneration": _HfExamplesInfo( + "deepseek-ai/DeepSeek-V4-Flash-Vision-Exp", + ), "Dots3NoteForCausalLM": _HfExamplesInfo( "dots-studio/dots3-note-prev", is_available_online=False, diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 395d7e3706..3566eb5341 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -132,6 +132,12 @@ def can_initialize( f"capability {capability.major}.{capability.minor}" ) + if model_arch == "DeepseekV4ForConditionalGeneration": + from vllm.platforms import current_platform + + if not current_platform.is_cuda(): + pytest.skip("Deepseek V4 is only supported on CUDA") + with ( patch.object(V1EngineCore, "_initialize_kv_caches", _initialize_kv_caches_v1), monkeypatch.context() as m, diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 25eeb3f6a0..1568fcc9bb 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -64,6 +64,11 @@ def test_registry_imports(model_arch): ): pytest.skip("HY V4 is only supported on CUDA") + if model_arch == "DeepseekV4ForConditionalGeneration" and ( + current_platform.is_cuda() + ): + pytest.skip("Deepseek V4 is only supported on CUDA") + # Ensure all model classes can be imported successfully model_cls = ModelRegistry._try_load_model_cls(model_arch) assert model_cls is not None diff --git a/tests/tokenizers_/test_deepseek_v4.py b/tests/tokenizers_/test_deepseek_v4.py index 6fd4be9fd8..a55aa57666 100644 --- a/tests/tokenizers_/test_deepseek_v4.py +++ b/tests/tokenizers_/test_deepseek_v4.py @@ -6,6 +6,7 @@ from pathlib import Path from types import SimpleNamespace import pytest +import torch from vllm.entrypoints.chat_utils import parse_chat_messages from vllm.renderers.registry import RENDERER_REGISTRY @@ -427,3 +428,60 @@ def test_deepseek_v4_encode_messages_rejects_invalid_arguments(kwargs): with pytest.raises(ValueError): encode_messages([{"role": "user", "content": "Hello"}], **kwargs) + + +def test_deepseek_v4_image_blocks_become_placeholders(): + prompt = _tokenizer().apply_chat_template( + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "first:"}, + {"type": "image_url", "image_url": {"url": "file:///a.png"}}, + {"type": "text", "text": "second:"}, + {"type": "image", "source": {"data": "AAAA"}}, + ], + } + ], + tokenize=False, + thinking=False, + ) + + assert "<|User|>first:<|deepseek_image|>second:<|deepseek_image|>" in prompt + + +def test_deepseek_v4_image_sentinel_ids_match_tokenizer(): + """The borrowed sentinel ids must line up with the reserved + ``<|place_holder_mm_span_XXXX|>`` tokens in the tokenizer.""" + from vllm.models.deepseek_v4.common.mm_preprocess import ( + IMAGE_SENTINEL_BASE_ID, + IMAGE_SENTINEL_TOKEN_NAMES, + image_sentinel_mask, + validate_image_sentinel_ids, + ) + + class FakeTokenizer: + def __init__(self, offset: int = 0) -> None: + self.offset = offset + + def convert_tokens_to_ids(self, token: str) -> int: + return ( + IMAGE_SENTINEL_BASE_ID + + self.offset + + (IMAGE_SENTINEL_TOKEN_NAMES.index(token)) + ) + + validate_image_sentinel_ids(FakeTokenizer()) # no raise + with pytest.raises(ValueError, match="sentinel"): + validate_image_sentinel_ids(FakeTokenizer(offset=1)) + + ids = torch.tensor( + [ + 1, + IMAGE_SENTINEL_BASE_ID, + IMAGE_SENTINEL_BASE_ID + 4, + IMAGE_SENTINEL_BASE_ID + 5, + 129264, + ] + ) + assert image_sentinel_mask(ids).tolist() == [False, True, True, False, False] diff --git a/tests/v1/attention/test_deepseek_v4_swa_visible.py b/tests/v1/attention/test_deepseek_v4_swa_visible.py new file mode 100644 index 0000000000..d18b7c671a --- /dev/null +++ b/tests/v1/attention/test_deepseek_v4_swa_visible.py @@ -0,0 +1,580 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for DeepSeek-V4 in-image bidirectional SWA visibility (vision variant). + +The reference semantics are a torch port of the official +``get_image_visible`` / ``get_window_topk_idxs_visible``: inside an image span +[span_start, span_end] (inclusive), a token at ``pos`` sees +``min(pos - span_start, max_image_tokens - 1)`` extra tokens to the left and +``min(span_end - pos, max_image_tokens)`` to the right; the window then starts +at ``max(pos - (window - 1) - max(left - (window - 1), 0), 0)`` and ends at +``pos + right`` (inclusive). +""" + +import pytest +import torch +from typing_extensions import TypedDict + +from tests.v1.attention.utils import create_vllm_config +from vllm.models.deepseek_v4.common.ops.cache_utils import ( + build_flashinfer_mixed_sparse_indices, + combine_topk_swa_indices, +) +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadataBuilder, + _compute_image_visibility_kernel, + _compute_swa_indices_and_lens_kernel, +) +from vllm.v1.kv_cache_interface import SlidingWindowMLASpec + +WINDOW = 8 +MAX_IMG = 6 +WIDTH = WINDOW + MAX_IMG +BLOCK_SIZE = 64 + + +def ref_left_right( + seq_lens: list[int], + query_lens: list[int], + spans_per_req: list[list[tuple[int, int]]], + max_image_tokens: int, +) -> tuple[list[int], list[int]]: + """Per-token (left, right) for the flattened decode-first token stream.""" + lefts: list[int] = [] + rights: list[int] = [] + for seq_len, query_len, spans in zip(seq_lens, query_lens, spans_per_req): + prefix_len = seq_len - query_len + for i in range(query_len): + pos = prefix_len + i + left = right = 0 + for span_start, span_end in spans: + if span_start <= pos <= span_end: + left = min(pos - span_start, max_image_tokens - 1) + right = min(span_end - pos, max_image_tokens) + lefts.append(left) + rights.append(right) + return lefts, rights + + +def ref_swa_bounds(pos: int, window: int, left: int, right: int) -> tuple[int, int]: + """Reference [start, end) window bounds for one query token.""" + left_add = max(left - (window - 1), 0) + start = max(pos - (window - 1) - left_add, 0) + return start, pos + right + 1 + + +def make_batch( + seq_lens: list[int], + query_lens: list[int], + device: torch.device, +): + """Build the kernel inputs for a decode-first batch.""" + num_reqs = len(seq_lens) + query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) + query_start_loc[1:] = torch.tensor( + query_lens, dtype=torch.int32, device=device + ).cumsum(0) + num_tokens = int(query_start_loc[-1]) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + token_to_req = torch.repeat_interleave( + torch.arange(num_reqs, dtype=torch.int32, device=device), + torch.tensor(query_lens, dtype=torch.int32, device=device), + ) + max_blocks = (max(seq_lens) + BLOCK_SIZE - 1) // BLOCK_SIZE + block_table = torch.arange( + num_reqs * max_blocks, dtype=torch.int32, device=device + ).view(num_reqs, max_blocks) + slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device) + return query_start_loc, seq_lens_t, token_to_req, slot_mapping, block_table + + +def ref_swa_slot_rows( + seq_lens: list[int], + query_lens: list[int], + spans_per_req: list[list[tuple[int, int]]], + block_table: torch.Tensor, + window: int, + max_image_tokens: int, + width: int, +) -> tuple[list[list[int]], list[int]]: + """Reference paged slot-id rows and lens for every token in the batch.""" + block_table_cpu = block_table.cpu() + lefts, rights = ref_left_right( + seq_lens, query_lens, spans_per_req, max_image_tokens + ) + rows: list[list[int]] = [] + lens: list[int] = [] + token = 0 + for req, (seq_len, query_len) in enumerate(zip(seq_lens, query_lens)): + prefix_len = seq_len - query_len + for i in range(query_len): + pos = prefix_len + i + start, end = ref_swa_bounds(pos, window, lefts[token], rights[token]) + row = [] + for p in range(start, end): + blk = int(block_table_cpu[req, p // BLOCK_SIZE]) + row.append(blk * BLOCK_SIZE + p % BLOCK_SIZE) + lens.append(len(row)) + row.extend([-1] * (width - len(row))) + rows.append(row) + token += 1 + return rows, lens + + +def run_swa_kernel( + seq_lens: list[int], + query_lens: list[int], + spans_per_req: list[list[tuple[int, int]]], + window: int = WINDOW, + max_image_tokens: int = MAX_IMG, + with_image: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + device = torch.device("cuda") + query_start_loc, seq_lens_t, token_to_req, slot_mapping, block_table = make_batch( + seq_lens, query_lens, device + ) + num_tokens = int(query_start_loc[-1]) + width = window + (max_image_tokens if with_image else 0) + swa_indices = torch.zeros(num_tokens, 1, width, dtype=torch.int32, device=device) + swa_lens = torch.zeros(num_tokens, dtype=torch.int32, device=device) + is_valid = slot_mapping >= 0 + + if with_image: + lefts, rights = ref_left_right( + seq_lens, query_lens, spans_per_req, max_image_tokens + ) + left_t = torch.tensor(lefts, dtype=torch.int32, device=device) + right_t = torch.tensor(rights, dtype=torch.int32, device=device) + else: + left_t = right_t = swa_lens # unused dummies + + _compute_swa_indices_and_lens_kernel[(num_tokens,)]( + swa_indices, + swa_indices.stride(0), + swa_lens, + window, + width, + left_t, + right_t, + query_start_loc, + seq_lens_t, + token_to_req, + is_valid, + block_table, + block_table.stride(0), + BLOCK_SIZE, + token_offset=0, + HAS_IMAGE=with_image, + TRITON_BLOCK_SIZE=1024, + ) + return swa_indices[:, 0], swa_lens + + +# seq_lens, query_lens, spans per request (positions are prompt-absolute, +# inclusive on both ends, matching the reference's sentinel-bracketed spans). +class _Case(TypedDict): + seq_lens: list[int] + query_lens: list[int] + spans: list[list[tuple[int, int]]] + + +CASES: list[_Case] = [ + # two image spans in one request, one pure-text request + { + "seq_lens": [30, 12], + "query_lens": [30, 12], + "spans": [[(4, 12), (20, 24)], []], + }, + # span larger than max_image_tokens, straddling the window boundary + { + "seq_lens": [26], + "query_lens": [26], + "spans": [[(2, 17)]], + }, + # chunked prefill: span fully inside the second chunk's query + { + "seq_lens": [40, 9], + "query_lens": [16, 9], + "spans": [[(30, 38)], []], + }, + # span ending exactly at the prompt end; tiny request + { + "seq_lens": [7, 15], + "query_lens": [7, 15], + "spans": [[(0, 6)], [(10, 14)]], + }, +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("case", CASES) +def test_swa_indices_kernel_with_image_spans(case): + rows, lens = ref_swa_slot_rows( + case["seq_lens"], + case["query_lens"], + case["spans"], + make_batch(case["seq_lens"], case["query_lens"], torch.device("cuda"))[4], + WINDOW, + MAX_IMG, + WIDTH, + ) + indices, actual_lens = run_swa_kernel( + case["seq_lens"], case["query_lens"], case["spans"], with_image=True + ) + assert actual_lens.cpu().tolist() == lens + assert indices.cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("case", CASES) +def test_swa_indices_kernel_without_image_unchanged(case): + """HAS_IMAGE=False must reproduce the plain causal sliding window.""" + rows, lens = ref_swa_slot_rows( + case["seq_lens"], + case["query_lens"], + [[] for _ in case["seq_lens"]], + make_batch(case["seq_lens"], case["query_lens"], torch.device("cuda"))[4], + WINDOW, + MAX_IMG, + WINDOW, + ) + indices, actual_lens = run_swa_kernel( + case["seq_lens"], case["query_lens"], case["spans"], with_image=False + ) + assert indices.shape[-1] == WINDOW + assert actual_lens.cpu().tolist() == lens + assert indices.cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_image_visibility_kernel(): + """The builder's visibility kernel must match get_image_visible.""" + device = torch.device("cuda") + seq_lens = [30, 12] + query_lens = [30, 12] + spans = [[(4, 12), (20, 24)], []] + query_start_loc, seq_lens_t, token_to_req, _, _ = make_batch( + seq_lens, query_lens, device + ) + num_tokens = int(query_start_loc[-1]) + + # CSR span layout: request -> contiguous [start, end) span list. + indptr = [0, 2, 2] + starts = [4, 20] + ends = [12, 24] + left = torch.zeros(num_tokens, dtype=torch.int32, device=device) + right = torch.zeros(num_tokens, dtype=torch.int32, device=device) + _compute_image_visibility_kernel[(num_tokens,)]( + left, + right, + torch.tensor(indptr, dtype=torch.int32, device=device), + torch.tensor(starts, dtype=torch.int32, device=device), + torch.tensor(ends, dtype=torch.int32, device=device), + query_start_loc, + seq_lens_t, + token_to_req, + MAX_IMG, + token_offset=0, + ) + ref_left, ref_right = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) + assert left.cpu().tolist() == ref_left + assert right.cpu().tolist() == ref_right + + +def combine_case( + compress_ratio: int, + topk: int, + seq_lens: list[int], + query_lens: list[int], + spans: list[list[tuple[int, int]]], + with_image: bool, +): + """Run combine_topk_swa_indices and return (indices, lens, expected).""" + device = torch.device("cuda") + num_reqs = len(seq_lens) + query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) + query_start_loc[1:] = torch.tensor( + query_lens, dtype=torch.int32, device=device + ).cumsum(0) + num_tokens = int(query_start_loc[-1]) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + gather_lens = torch.tensor( + [q + min(s - q, WINDOW - 1) for s, q in zip(seq_lens, query_lens)], + dtype=torch.int32, + device=device, + ) + N = (max(seq_lens) + compress_ratio - 1) // compress_ratio + M = N + int(gather_lens.max()) + 8 + gen = torch.Generator(device="cpu").manual_seed(0) + topk_indices = torch.randint( + 0, 4096, (num_tokens, max(topk, 1)), generator=gen, dtype=torch.int32 + ).to(device) + topk_indices = topk_indices[:, : max(topk, 1)] + + if with_image: + lefts, rights = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) + left_t = torch.tensor(lefts, dtype=torch.int32, device=device) + right_t = torch.tensor(rights, dtype=torch.int32, device=device) + else: + left_t = right_t = None + + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices, + query_start_loc, + seq_lens_t, + gather_lens, + WINDOW, + compress_ratio, + topk, + M, + N, + left_visible=left_t, + right_visible=right_t, + max_image_tokens=MAX_IMG, + ) + + # Reference rows. + lefts, rights = ref_left_right( + seq_lens, + query_lens, + spans if with_image else [[] for _ in seq_lens], + MAX_IMG, + ) + topk_cpu = topk_indices.cpu() + width = WINDOW + MAX_IMG + combined_topk = (topk + width + 127) // 128 * 128 + rows = [] + lens = [] + token = 0 + for b, (seq_len, query_len) in enumerate(zip(seq_lens, query_lens)): + prefix_len = seq_len - query_len + gather_start = seq_len - int(gather_lens[b]) + for i in range(query_len): + pos = prefix_len + i + topk_len = min((pos + 1) // compress_ratio, topk) + start, end = ref_swa_bounds(pos, WINDOW, lefts[token], rights[token]) + swa_len = end - start + row = [-1] * combined_topk + for j in range(topk_len): + row[j] = int(topk_cpu[token, j]) + M * b + for j in range(swa_len): + row[topk_len + j] = M * b + N + start + j - gather_start + rows.append(row) + lens.append(topk_len + swa_len) + token += 1 + return combined_indices, combined_lens, rows, lens + + +COMBINE_CASES = [ + dict(compress_ratio=1, topk=0), # SWA-only layer + dict(compress_ratio=4, topk=16), # C4A layer +] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("cfg", COMBINE_CASES) +def test_combine_topk_swa_with_image_spans(cfg): + case = CASES[0] + indices, lens, rows, exp_lens = combine_case( + cfg["compress_ratio"], + cfg["topk"], + case["seq_lens"], + case["query_lens"], + case["spans"], + with_image=True, + ) + assert lens.cpu().tolist() == exp_lens + assert indices.cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +@pytest.mark.parametrize("cfg", COMBINE_CASES) +def test_combine_topk_swa_without_image_unchanged(cfg): + """left_visible=None must reproduce the plain causal combined indices.""" + case = CASES[0] + indices, lens, rows, exp_lens = combine_case( + cfg["compress_ratio"], + cfg["topk"], + case["seq_lens"], + case["query_lens"], + case["spans"], + with_image=False, + ) + assert lens.cpu().tolist() == exp_lens + assert indices.cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_flashinfer_mixed_sparse_indices_with_image_spans(): + device = torch.device("cuda") + # 1 decode token (req 0) + two prefill requests (reqs 1, 2). + seq_lens = [20, 30, 9] + query_lens = [1, 30, 9] + spans = [[], [(4, 12)], [(0, 6)]] + query_start_loc = torch.tensor([0, 1, 31, 40], dtype=torch.int32, device=device) + seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device) + token_to_req = torch.tensor( + [0] + [1] * 30 + [2] * 9, dtype=torch.int32, device=device + ) + block_table = torch.arange(3, dtype=torch.int32, device=device).view(3, 1) + decode_swa = torch.tensor( + [[100, 101, 102, 103, 104, 105, 106, 107]], dtype=torch.int32, device=device + ) + prefill_topk = torch.zeros(39, 0, dtype=torch.int32, device=device) + + lefts, rights = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) + left_t = torch.tensor(lefts, dtype=torch.int32, device=device) + right_t = torch.tensor(rights, dtype=torch.int32, device=device) + + sparse_indices, sparse_lens = build_flashinfer_mixed_sparse_indices( + decode_swa_indices=decode_swa, + decode_compressed_indices=None, + decode_compressed_topk_lens=None, + prefill_topk_indices=prefill_topk, + query_start_loc=query_start_loc, + seq_lens=seq_lens_t, + token_to_req_indices=token_to_req, + swa_block_table=block_table, + swa_block_size=BLOCK_SIZE, + compressed_block_table=None, + compressed_block_size=BLOCK_SIZE, + window_size=WINDOW, + compress_ratio=1, + topk=0, + prefill_left_visible=left_t, + prefill_right_visible=right_t, + max_image_tokens=MAX_IMG, + ) + + swa_total = WINDOW + MAX_IMG + assert sparse_indices.shape == (40, swa_total) + assert sparse_lens.cpu().tolist() == [swa_total] * 40 + # Decode row: slots copied, image-extension columns padded with -1. + decode_row = sparse_indices[0].cpu().tolist() + assert decode_row[:WINDOW] == list(range(100, 108)) + assert decode_row[WINDOW:] == [-1] * MAX_IMG + # Prefill rows: paged slot ids over the widened window. + exp_rows, _ = ref_swa_slot_rows( + seq_lens, query_lens, spans, block_table, WINDOW, MAX_IMG, swa_total + ) + actual = sparse_indices[1:].cpu().tolist() + assert actual == exp_rows[1:] + + +def make_builder(vision: bool) -> DeepseekSparseSWAMetadataBuilder: + overrides: dict = {"sliding_window": WINDOW} + if vision: + overrides.update(vision_n_layers=2, vision_max_n_token=MAX_IMG) + vllm_config = create_vllm_config( + max_model_len=4096, + max_num_batched_tokens=64, + max_num_seqs=8, + hf_config_override=overrides, + ) + spec = SlidingWindowMLASpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=512, + dtype=torch.bfloat16, + sliding_window=WINDOW, + cache_dtype_str="auto", + model_version="deepseek_v4", + ) + return DeepseekSparseSWAMetadataBuilder( + kv_cache_spec=spec, + layer_names=["layer0"], + vllm_config=vllm_config, + device=torch.device("cuda"), + ) + + +def build_metadata( + builder: DeepseekSparseSWAMetadataBuilder, + seq_lens: list[int], + query_lens: list[int], + mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None, +): + device = torch.device("cuda") + query_start_loc, seq_lens_t, _, slot_mapping, block_table = make_batch( + seq_lens, query_lens, device + ) + return builder.build( + 0, + CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc.cpu(), + seq_lens=seq_lens_t, + seq_lens_cpu_upper_bound=seq_lens_t.cpu(), + num_reqs=len(seq_lens), + num_actual_tokens=int(query_start_loc[-1]), + max_query_len=max(query_lens), + max_seq_len=max(seq_lens), + block_table_tensor=block_table, + slot_mapping=slot_mapping, + causal=True, + mm_req_doc_ranges=mm_req_doc_ranges, + ), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_builder_in_image_visibility(): + seq_lens = [30, 12] + query_lens = [30, 12] + spans = [[(4, 12), (20, 24)], []] + builder = make_builder(vision=True) + md = build_metadata(builder, seq_lens, query_lens, {0: spans[0], 1: spans[1]}) + assert md.num_prefills == 2 + assert md.num_decode_tokens == 0 + assert md.prefill_swa_indices.shape[-1] == WIDTH + assert md.prefill_left_visible is not None + assert md.prefill_right_visible is not None + + ref_left, ref_right = ref_left_right(seq_lens, query_lens, spans, MAX_IMG) + assert md.prefill_left_visible.cpu().tolist() == ref_left + assert md.prefill_right_visible.cpu().tolist() == ref_right + + _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) + rows, lens = ref_swa_slot_rows( + seq_lens, query_lens, spans, block_table, WINDOW, MAX_IMG, WIDTH + ) + assert md.prefill_swa_lens.cpu().tolist() == lens + assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_builder_no_image_spans_fast_path(): + """Vision model, image-free batch: no visibility tensors, plain window.""" + seq_lens = [30, 12] + query_lens = [30, 12] + builder = make_builder(vision=True) + md = build_metadata(builder, seq_lens, query_lens, {0: [], 1: []}) + assert md.prefill_left_visible is None + assert md.prefill_right_visible is None + + _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) + rows, lens = ref_swa_slot_rows( + seq_lens, query_lens, [[], []], block_table, WINDOW, MAX_IMG, WIDTH + ) + assert md.prefill_swa_lens.cpu().tolist() == lens + assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_builder_text_model_unchanged(): + """Text-only model: buffers stay window-sized and spans are ignored.""" + seq_lens = [30, 12] + query_lens = [30, 12] + builder = make_builder(vision=False) + assert builder.max_image_tokens == 0 + md = build_metadata(builder, seq_lens, query_lens, None) + assert md.prefill_swa_indices.shape[-1] == WINDOW + assert md.prefill_left_visible is None + + _, _, _, _, block_table = make_batch(seq_lens, query_lens, torch.device("cuda")) + rows, lens = ref_swa_slot_rows( + seq_lens, query_lens, [[], []], block_table, WINDOW, MAX_IMG, WINDOW + ) + assert md.prefill_swa_lens.cpu().tolist() == lens + assert md.prefill_swa_indices[:, 0].cpu().tolist() == rows diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 04b8b829f5..ecc3cf3c9f 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2501,6 +2501,8 @@ def topk_hash_softplus_sqrt( input_tokens: torch.Tensor | None = None, hash_indices_table: torch.Tensor | None = None, is_padding: torch.Tensor | None = None, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ) -> None: torch.ops._moe_C.topk_softplus_sqrt( topk_weights, @@ -2513,6 +2515,8 @@ def topk_hash_softplus_sqrt( input_tokens, hash_indices_table, is_padding, + bias_vl, + image_sentinel_lo, ) diff --git a/vllm/config/model.py b/vllm/config/model.py index cfb2ce3e63..1598087028 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -697,7 +697,10 @@ class ModelConfig: self.tokenizer_mode = "kimi_k3" elif arch == "DeepseekV32ForCausalLM": self.tokenizer_mode = "deepseek_v32" - elif arch == "DeepseekV4ForCausalLM": + elif arch in ( + "DeepseekV4ForCausalLM", + "DeepseekV4ForConditionalGeneration", + ): self.tokenizer_mode = "deepseek_v4" elif arch in ("InklingForCausalLM", "InklingForConditionalGeneration"): self.tokenizer_mode = "inkling" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5181bbc53b..e4c223daf6 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -77,6 +77,7 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( "DeepseekV32MTPModel", "DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", + "DeepseekV4ForConditionalGeneration", "DeepSeekV4MTPModel", "Dots3NoteForCausalLM", "Dots3NoteMTPModel", diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index d9fd1385a9..4b633dd99b 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -131,6 +131,8 @@ def FusedMoEFactory( apply_routed_scale_to_output: bool = False, zero_expert_type: str | None = None, hash_indices_table: torch.Tensor | None = None, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, runner_cls: type[MoERunner] | None = None, runner_args: dict[str, Any] | None = None, routed_experts_cls: type[RoutedExperts] | None = None, @@ -199,6 +201,9 @@ def FusedMoEFactory( output instead of topk_weights zero_expert_type: Type of zero expert handling hash_indices_table: Hash table for expert indices + bias_vl: Vision routing bias for image tokens (Deepseek V4) + image_sentinel_lo: First of five consecutive in-vocab image sentinel + ids (0 = vision routing disabled) runner_cls: Custom MoERunner class (None = use default MoERunner) runner_args: Additional arguments for runner constructor routed_experts_cls: Custom RoutedExperts class (None = use default) @@ -313,6 +318,8 @@ def FusedMoEFactory( zero_expert_type=zero_expert_type, num_logical_experts=logical_num_experts, hash_indices_table=hash_indices_table, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) if params_dtype is None: diff --git a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py index dd861d4ebc..dbd29f3d08 100644 --- a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py +++ b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py @@ -44,8 +44,12 @@ if current_platform.is_cuda(): topk_weights_ptr, topk_ids_ptr, routed_scaling_factor, + input_ids_ptr, + bias_vl_ptr, + image_sentinel_lo, NUM_EXPERTS: tl.constexpr, BLOCK_N: tl.constexpr, + HAS_VL: tl.constexpr, launch_pdl: tl.constexpr, ): row = tl.program_id(0) @@ -54,6 +58,19 @@ if current_platform.is_cuda(): bias = tl.load( correction_bias_ptr + expert_offsets, mask=expert_mask, other=0.0 ).to(tl.float32) + if HAS_VL: + # Image tokens carry five consecutive in-vocab sentinel ids + # starting at image_sentinel_lo and use bias_vl for expert + # selection instead of the regular correction bias. Ids above the + # sentinel block are regular special tokens and must not match. + token_id = tl.load(input_ids_ptr + row).to(tl.int64) + bias_vl = tl.load( + bias_vl_ptr + expert_offsets, mask=expert_mask, other=0.0 + ).to(tl.float32) + is_image = (token_id >= image_sentinel_lo) & ( + token_id < image_sentinel_lo + 5 + ) + bias = tl.where(is_image, bias_vl, bias) if launch_pdl: tl.extra.cuda.gdc_wait() @@ -101,8 +118,17 @@ def dsv4_topk( correction_bias: torch.Tensor, indices_dtype: torch.dtype, routed_scaling_factor: float, + input_ids: torch.Tensor | None = None, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens, num_experts = gating_output.shape + has_vl = bias_vl is not None and image_sentinel_lo > 0 + if bias_vl is not None: + assert input_ids is not None, "bias_vl routing requires input_ids" + assert bias_vl.dtype == torch.float32 and bias_vl.is_contiguous() + assert bias_vl.shape == (num_experts,) + assert input_ids.is_contiguous() shape = (num_tokens, _TOPK) topk_weights = gating_output.new_empty(shape, dtype=torch.float32) topk_ids = gating_output.new_empty(shape, dtype=indices_dtype) @@ -113,8 +139,12 @@ def dsv4_topk( topk_weights, topk_ids, routed_scaling_factor, + input_ids, + bias_vl, + image_sentinel_lo, NUM_EXPERTS=num_experts, BLOCK_N=triton.next_power_of_2(num_experts), + HAS_VL=has_vl, num_warps=1, launch_pdl=current_platform.is_arch_support_pdl(), ) diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index f1c2636156..257ede27d5 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -81,6 +81,8 @@ def vllm_topk_softplus_sqrt( input_tokens: torch.Tensor | None = None, hash_indices_table: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ) -> tuple[torch.Tensor, ...]: ops.topk_hash_softplus_sqrt( topk_weights, @@ -93,6 +95,8 @@ def vllm_topk_softplus_sqrt( input_tokens, hash_indices_table, is_padding=_get_padding_mask(topk_indices.shape[0]), + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) return topk_weights, topk_indices @@ -122,6 +126,8 @@ def fused_topk_bias( input_tokens: torch.Tensor | None = None, hash_indices_table: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ): if ( input_tokens is not None @@ -130,18 +136,28 @@ def fused_topk_bias( ): input_tokens = input_tokens.to(dtype=hash_indices_table.dtype) + if bias_vl is not None: + # Image tokens carry five consecutive in-vocab sentinel ids starting + # at image_sentinel_lo and select experts with bias_vl instead of the + # regular routing path. image_sentinel_lo == 0 disables this. + assert input_tokens is not None, "bias_vl routing requires input_tokens" + if not rocm_aiter_ops.is_fused_moe_enabled(): assert hidden_states.size(0) == gating_output.size(0), ( "Number of tokens mismatch" ) output_indices_dtype = torch.int32 if indices_type is None else indices_type - if scoring_func == "sqrtsoftplus" and can_use_dsv4_topk( - gating_output, - e_score_correction_bias, - topk, - renormalize, - output_indices_dtype, + if ( + scoring_func == "sqrtsoftplus" + and hash_indices_table is None + and can_use_dsv4_topk( + gating_output, + e_score_correction_bias, + topk, + renormalize, + output_indices_dtype, + ) ): assert e_score_correction_bias is not None return dsv4_topk( @@ -149,6 +165,9 @@ def fused_topk_bias( e_score_correction_bias, output_indices_dtype, routed_scaling_factor, + input_ids=input_tokens, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) M, _ = hidden_states.size() @@ -200,6 +219,8 @@ def fused_topk_bias( input_tokens, hash_indices_table, routed_scaling_factor, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) else: raise ValueError(f"Unsupported scoring function: {scoring_func}") @@ -255,6 +276,8 @@ def fused_topk_bias( input_tokens, hash_indices_table, routed_scaling_factor, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) n_routed_experts = gating_output.shape[-1] @@ -270,10 +293,39 @@ def fused_topk_bias( ) + e_score_correction_bias.unsqueeze(0) else: scores_for_choice = scores.view(-1, n_routed_experts) + image_mask = None + if bias_vl is not None and image_sentinel_lo > 0: + # Image tokens (five consecutive sentinel ids starting at + # image_sentinel_lo) select experts with bias_vl instead of + # e_score_correction_bias / the hash table. Ids above the sentinel + # block are regular special tokens and must not match. + assert input_tokens is not None, "bias_vl routing requires input_tokens" + image_mask = ( + (input_tokens >= image_sentinel_lo) & (input_tokens < image_sentinel_lo + 5) + ).unsqueeze(-1) + text_bias = ( + e_score_correction_bias + if e_score_correction_bias is not None + else torch.zeros_like(bias_vl) + ) + row_bias = torch.where(image_mask, bias_vl, text_bias) + scores_for_choice = scores.view(-1, n_routed_experts) + row_bias # For batch invariance, use sorted=True to ensure deterministic expert selection if hash_indices_table is not None: assert input_tokens is not None - topk_indices = hash_indices_table[input_tokens] + if image_mask is None: + topk_indices = hash_indices_table[input_tokens] + else: + # Clamp sentinel rows (overwritten below) to keep the table + # lookup well-defined. + safe_ids = torch.where(image_mask.squeeze(-1), 0, input_tokens).to( + hash_indices_table.dtype + ) + topk_indices = hash_indices_table[safe_ids] + vl_indices = torch.topk(scores_for_choice, k=topk, dim=-1)[1] + topk_indices = torch.where( + image_mask, vl_indices.to(topk_indices.dtype), topk_indices + ) else: use_sorted = envs.VLLM_BATCH_INVARIANT topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[ @@ -306,6 +358,8 @@ class FusedTopKBiasRouter(BaseRouter): hash_indices_table: torch.Tensor | None = None, num_fused_shared_experts: int = 0, shared_expert_weight: float = 1.0, + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ): super().__init__( top_k=top_k, @@ -318,6 +372,11 @@ class FusedTopKBiasRouter(BaseRouter): self.routed_scaling_factor = routed_scaling_factor self.scoring_func = scoring_func self._hash_indices_table = hash_indices_table + # Vision bias: image sentinel tokens (five consecutive in-vocab ids + # starting at image_sentinel_lo) select experts with bias_vl instead + # of e_score_correction_bias / the hash table. + self.bias_vl = bias_vl + self.image_sentinel_lo = image_sentinel_lo # Fused shared experts: append constant slots (ids immediately after # the routed experts, [global, global+n)) routed to by every token at # ``shared_expert_weight``, AFTER the routed top-k is renormalized. @@ -357,6 +416,8 @@ class FusedTopKBiasRouter(BaseRouter): input_tokens=input_ids, hash_indices_table=self._hash_indices_table, routed_scaling_factor=self.routed_scaling_factor, + bias_vl=self.bias_vl.data if self.bias_vl is not None else None, + image_sentinel_lo=self.image_sentinel_lo, ) if self.num_fused_shared_experts > 0: diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index ff4874b20e..d80fa840f8 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -60,6 +60,9 @@ def create_fused_moe_router( zero_expert_type: str | None = None, num_logical_experts: int | None = None, hash_indices_table: torch.Tensor | None = None, + # Deepseek V4 vision routing bias parameters + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ) -> FusedMoERouter: """ Factory function to create the appropriate FusedMoERouter subclass based on @@ -211,6 +214,8 @@ def create_fused_moe_router( hash_indices_table=hash_indices_table, num_fused_shared_experts=num_fused_shared_experts, shared_expert_weight=shared_expert_weight, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) if ( diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 263c8a4035..485d15354d 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -1002,6 +1002,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = { "ColBERTJinaRobertaModel": JinaRobertaModelConfig, "ColQwen3_5": ColQwen3_5Config, "DeepseekV4ForCausalLM": DeepseekV4ForCausalLMConfig, + "DeepseekV4ForConditionalGeneration": DeepseekV4ForCausalLMConfig, "DeepseekV32ForCausalLM": DeepseekV32ForCausalLM, "DiffusionGemmaForBlockDiffusion": DiffusionGemmaModelForBlockDiffusionConfig, # noqa: E501 "Ernie4_5_VLMoeForConditionalGeneration": Ernie4_5_VLMoeForConditionalGenerationConfig, # noqa: E501 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 40f4a27d28..5c1f3da254 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -362,6 +362,10 @@ _MULTIMODAL_MODELS = { "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"), "DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"), "DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"), + "DeepseekV4ForConditionalGeneration": ( + "vllm.models.deepseek_v4", + "DeepseekV4ForConditionalGeneration", + ), "Dots3NoteForCausalLM": ( "vllm.models.dots3_note", "Dots3NoteForCausalLM", diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index c7dd948f79..d71bddae33 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -20,7 +20,13 @@ if current_platform.is_rocm(): ) from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP + from .vl_stub import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) elif current_platform.is_xpu(): + from .vl_stub import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] @@ -30,10 +36,14 @@ else: ) from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment] + from .nvidia.vl_model import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) __all__ = [ "DSparkDeepseekV4ForCausalLM", "DeepSeekV4MTP", "DeepseekV4FP8Config", "DeepseekV4ForCausalLM", + "DeepseekV4ForConditionalGeneration", ] diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index e53cc7c2a8..ce9fe06a2f 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -205,6 +205,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): self.n_groups = config.o_groups self.n_local_groups = self.n_groups // tp_size self.window_size = config.sliding_window + # Vision variant: image spans are visible bidirectionally, widening + # prefill SWA index rows by up to max_image_tokens columns. + self.max_image_tokens = ( + getattr(config, "vision_max_n_token", 0) + if getattr(config, "vision_n_layers", 0) > 0 + else 0 + ) # NOTE(zyongye) Compress ratio can't be 0 # we do this for because MTP layer is not included # in the compress ratio list @@ -363,6 +370,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup( window_size=self.window_size, block_size=self.swa_cache_layer.block_size, + max_image_tokens=self.max_image_tokens, ) if self.compress_ratio > 1: diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py new file mode 100644 index 0000000000..919146c551 --- /dev/null +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multimodal preprocessing for the DeepSeek-V4 vision variants +(DeepSeek-V4-Flash-Vision-Exp). + +The image transform and sentinel-block construction are ported from the +official repository's ``image_processor.py`` so that token counts bit-match +the reference. Each ``<|deepseek_image|>`` placeholder in the prompt expands +to a variable-length block of sentinel tokens; only positions with +``type == IMAGE`` receive vision embeddings, the other sentinels are looked +up from learned embedding vectors in the model. + +Unlike the reference (which uses out-of-vocab ids ``vocab_size + type``), +the sentinel block borrows five consecutive reserved tokenizer tokens +(``<|place_holder_mm_span_0431|>`` .. ``_0435|>``): they are special tokens +the tokenizer never emits from plain text, so the ids stay in-vocabulary and +work with stock token validation, logprobs and detokenization. The ids are +pure markers — every sentinel position's embedding is overwritten with +vision/sentinel vectors before the decoder layers see it, exactly like the +reference's out-of-vocab scheme. +""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import numpy as np +import torch +from PIL import Image, ImageOps +from transformers import BatchFeature +from typing_extensions import assert_never + +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import ImageSize, MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.multimodal.processing.processor import ( + MultiModalPromptUpdates, + MultiModalPromptUpdatesApplyResult, + PlaceholderFeaturesInfo, + UpdateMode, + _plan_prompt_updates, +) +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config + +IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5) +COMPRESS_PAD_TO = 4 + +IMAGE_PLACEHOLDER = "<|deepseek_image|>" + +# Sentinel roles borrow five consecutive ``<|place_holder_mm_span_XXXX|>`` +# tokens (reserved special tokens, never emitted from plain text). The order +# must match IMAGE_START..IMAGE_END above so that +# ``id == IMAGE_SENTINEL_BASE_ID + type``. +IMAGE_SENTINEL_BASE_ID = 129257 +IMAGE_SENTINEL_TOKEN_NAMES = ( + "<|place_holder_mm_span_0431|>", # IMAGE_START + "<|place_holder_mm_span_0432|>", # IMAGE_PAD + "<|place_holder_mm_span_0433|>", # IMAGE + "<|place_holder_mm_span_0434|>", # IMAGE_NEW_LINE + "<|place_holder_mm_span_0435|>", # IMAGE_END +) + + +def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor: + """Boolean mask for image-block sentinel positions (in-vocab ids).""" + return (token_ids >= IMAGE_SENTINEL_BASE_ID) & ( + token_ids < IMAGE_SENTINEL_BASE_ID + len(IMAGE_SENTINEL_TOKEN_NAMES) + ) + + +def validate_image_sentinel_ids(tokenizer) -> None: + """Check the borrowed sentinel ids against the tokenizer.""" + for i, name in enumerate(IMAGE_SENTINEL_TOKEN_NAMES): + token_id = tokenizer.convert_tokens_to_ids(name) + if token_id != IMAGE_SENTINEL_BASE_ID + i: + raise ValueError( + f"Image sentinel token {name!r} has id {token_id}, expected " + f"{IMAGE_SENTINEL_BASE_ID + i}; the DeepSeek-V4 vision " + "sentinel block requires these consecutive reserved ids." + ) + + +def grid_tokens(best_height, best_width, patch_size, downsample_ratio): + """Number of LLM tokens the aligner grid occupies (N-layout, including + row/align padding).""" + n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio) + n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio) + num_tokens = n_llm_h * (n_llm_w + 1) + 2 + if n_llm_h % 2 == 1: + num_tokens += n_llm_w + 1 + num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 + return n_llm_h, n_llm_w, num_tokens + + +def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): + r = height / width + max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 + max_h_float = max_w_float * r + if max_w_float < 1.0: + max_w = 1 + max_h = (max_n_token - 2) // (max_w + 1) + if max_h % 2 == 1: + max_h -= 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + elif max_h_float < 2.0: + max_h = 2 + max_w = ((max_n_token - 2) // max_h) - 1 + assert max_w > 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + else: + max_w = math.floor(max_w_float) + max_h = math.floor(max_h_float) + if max_h % 2 == 1: + max_h -= 1 + beta = min( + max_w * patch_size * downsample_ratio / width, + max_h * patch_size * downsample_ratio / height, + ) + best_width = math.floor(width * beta / patch_size) * patch_size + best_height = math.floor(height * beta / patch_size) * patch_size + n_llm_h, n_llm_w, num_tokens = grid_tokens( + best_height, best_width, patch_size, downsample_ratio + ) + return n_llm_h, n_llm_w, best_height, best_width, num_tokens + + +def safe_resize( + height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token +): + max_n_token -= COMPRESS_PAD_TO - 1 + n_llm_h, n_llm_w, num_tokens = grid_tokens( + best_height, best_width, patch_size, downsample_ratio + ) + budget = max_n_token + while num_tokens > max_n_token: + n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio( + height, width, patch_size, downsample_ratio, budget + ) + budget -= 1 + return n_llm_h, n_llm_w, best_height, best_width + + +def load_image( + image: Image.Image, + *, + patch_size: int, + downsample_ratio: int, + max_n_token: int, + min_pixels: int, + max_wh_ratio: float | None, +): + """Transform one PIL image into ViT patches. + + Same math as the reference ``load_image``, except the image is already + decoded (vLLM supplies PIL images instead of a record dict). + """ + p = patch_size + image = image.convert("RGB") + width, height = image.size + if max_wh_ratio is not None and width > height * max_wh_ratio: + width = height * max_wh_ratio + if 0 < width * height < min_pixels: + ratio = (min_pixels / (width * height)) ** 0.5 + width = int(width * ratio) + height = int(height * ratio) + best_width = math.ceil(width / p) * p + best_height = math.ceil(height / p) * p + n_llm_h, n_llm_w, best_height, best_width = safe_resize( + height, width, best_height, best_width, p, downsample_ratio, max_n_token + ) + n_vit_h, n_vit_w = best_height // p, best_width // p + if max_wh_ratio is not None and image.width >= max_wh_ratio * image.height: + image = image.resize((best_width, best_height)) + else: + image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127)) + x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255 + x = ((x - 0.5) / 0.5).to(torch.bfloat16) + patches = ( + x.reshape(3, n_vit_h, p, n_vit_w, p) + .permute(1, 3, 0, 2, 4) + .reshape(n_vit_h * n_vit_w, 3, p, p) + ) + return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w + + +def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int): + """Builds the N-layout token types (final order) and the aligner-row order + for IMAGE slots.""" + compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO + pad_h = n_llm_h % 2 + rows = n_llm_h + pad_h + row_len = n_llm_w + 1 + pad_last = rows // 2 * row_len % 2 * 2 + types = torch.tensor( + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + + [IMAGE_PAD] * (row_len * pad_h), + dtype=torch.int64, + ) + order = torch.arange(rows * row_len).view(rows // 2, 2, row_len) + order = order.transpose(1, 2).reshape(-1) + image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64) + image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange( + n_llm_h * n_llm_w + ).view(n_llm_h, n_llm_w) + perm = image_idx[order] + perm = perm[perm >= 0] + types = torch.cat( + [ + torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_START]), + types[order], + torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_END]), + ] + ) + return types, perm + + +def build_image_block_pad_free(n_llm_h: int, n_llm_w: int): + """``build_image_block`` without the position-dependent compressor pad. + + ``start_pos`` is chosen so that ``compress_pad == 0``; the pad is instead + prepended when the block is spliced into the final prompt, where its + position is known. + """ + return build_image_block(n_llm_h, n_llm_w, COMPRESS_PAD_TO - 1) + + +class DeepseekV4VLImageProcessor: + """Per-image transform (the PIL-input equivalent of the reference + ``load_image``).""" + + def __init__(self, config: DeepseekV4Config) -> None: + super().__init__() + self.patch_size = config.vision_patch_size + self.downsample_ratio = config.vision_downsample_ratio + self.max_n_token = config.vision_max_n_token + self.min_pixels = config.vision_min_pixels + self.max_wh_ratio = config.vision_max_wh_ratio + + def __call__(self, image: Image.Image): + return load_image( + image, + patch_size=self.patch_size, + downsample_ratio=self.downsample_ratio, + max_n_token=self.max_n_token, + min_pixels=self.min_pixels, + max_wh_ratio=self.max_wh_ratio, + ) + + +class DeepseekV4VLProcessor: + """Minimal stand-in for the HF processor of DeepSeek-V4 vision models. + + The official repository ships image preprocessing as plain functions in + ``image_processor.py`` (no ``auto_map`` processor), so this class wraps + their ports directly and the model loads without ``--trust-remote-code``. + + ``__call__`` returns a ``BatchFeature`` with one entry per image + (flattened across images): + + - ``patches``: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` bf16 ViT patches. + - ``vit_grid``: ``(num_images, 2)`` int64 ``[n_vit_h, n_vit_w]``. + - ``llm_grid``: ``(num_images, 2)`` int64 ``[n_llm_h, n_llm_w]``. + - ``perm``: concatenated per-image ``(n_llm_h * n_llm_w,)`` int64 index + selecting aligner outputs into the final N-layout order. + - ``types``: concatenated per-image pad-free sentinel block types; + ``block ids = IMAGE_SENTINEL_BASE_ID + types``. + """ + + def __init__(self, config: DeepseekV4Config) -> None: + super().__init__() + self.config = config + self.image_processor = DeepseekV4VLImageProcessor(config) + + def __call__( + self, + text: str | None = None, + images: Sequence[Image.Image] | None = None, + return_tensors: str | None = None, + **kwargs: Any, + ) -> BatchFeature: + patches_list = [] + vit_grid = [] + llm_grid = [] + perm_list = [] + types_list = [] + for image in images or []: + patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = self.image_processor(image) + types, perm = build_image_block_pad_free(n_llm_h, n_llm_w) + patches_list.append(patches) + vit_grid.append((n_vit_h, n_vit_w)) + llm_grid.append((n_llm_h, n_llm_w)) + perm_list.append(perm) + types_list.append(types) + + if not patches_list: + return BatchFeature({}) + + return BatchFeature( + { + "patches": torch.cat(patches_list), + "vit_grid": torch.tensor(vit_grid, dtype=torch.int64), + "llm_grid": torch.tensor(llm_grid, dtype=torch.int64), + "perm": torch.cat(perm_list), + "types": torch.cat(types_list), + } + ) + + +class DeepseekV4VLProcessingInfo(BaseProcessingInfo): + def get_hf_config(self) -> DeepseekV4Config: + return self.ctx.get_hf_config(DeepseekV4Config) + + def get_hf_processor(self, **kwargs: object) -> DeepseekV4VLProcessor: + if kwargs: + raise ValueError(f"Unexpected processor kwargs: {sorted(kwargs)}") + return DeepseekV4VLProcessor(self.get_hf_config()) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + # ``safe_resize`` reserves COMPRESS_PAD_TO - 1 tokens of + # vision_max_n_token for the compressor-alignment pad, so the full + # sentinel block is bounded by vision_max_n_token; the margin is kept + # in case that reservation changes. + return { + "image": self.get_hf_config().vision_max_n_token + COMPRESS_PAD_TO - 1, + } + + def get_image_placeholder_token_id(self) -> int: + token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER) + if token_id is None: + raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}") + return token_id + + def get_image_size_with_most_features(self) -> ImageSize: + hf_config = self.get_hf_config() + patch_size = hf_config.vision_patch_size + downsample_ratio = hf_config.vision_downsample_ratio + # A square maximizes the ViT patch count (area) within the token + # budget; solve the budget-derived size directly to keep the dummy + # image small. + budget = hf_config.vision_max_n_token - (COMPRESS_PAD_TO - 1) + side = budget * patch_size * downsample_ratio + _, _, best_h, best_w, _ = solve_resize_ratio( + side, side, patch_size, downsample_ratio, budget + ) + return ImageSize(width=best_w, height=best_h) + + +class DeepseekV4VLDummyInputsBuilder( + BaseDummyInputsBuilder[DeepseekV4VLProcessingInfo] +): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return IMAGE_PLACEHOLDER * mm_counts.get("image", 0) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + size = self.info.get_image_size_with_most_features() + return { + "image": self._get_dummy_images( + width=size.width, + height=size.height, + num_images=mm_counts.get("image", 0), + overrides=cast(ImageDummyOptions | None, mm_options.get("image")), + ), + } + + +class DeepseekV4VLMultiModalProcessor( + BaseMultiModalProcessor[DeepseekV4VLProcessingInfo] +): + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + vit_grid = hf_inputs.get("vit_grid") + llm_grid = hf_inputs.get("llm_grid") + + if vit_grid is None or llm_grid is None: + empty = torch.empty(0, dtype=torch.long) + patch_sizes = perm_sizes = types_sizes = empty + else: + patch_sizes = vit_grid.prod(-1) + perm_sizes = llm_grid.prod(-1) + n_llm_h, n_llm_w = llm_grid[:, 0], llm_grid[:, 1] + # Pad-free block length; same formula as ``grid_tokens`` given + # the LLM grid. + types_sizes = ( + n_llm_h * (n_llm_w + 1) + + 2 + + (n_llm_h % 2) * (n_llm_w + 1) + + (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 + ) + + return { + "patches": MultiModalFieldConfig.flat_from_sizes("image", patch_sizes), + "vit_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), + "llm_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), + "perm": MultiModalFieldConfig.flat_from_sizes("image", perm_sizes), + "types": MultiModalFieldConfig.flat_from_sizes( + "image", types_sizes, keep_on_cpu=True + ), + } + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + image_token_id = self.info.get_image_placeholder_token_id() + validate_image_sentinel_ids(self.info.get_tokenizer()) + image_embed_id = IMAGE_SENTINEL_BASE_ID + IMAGE + + def get_image_replacement(item_idx: int) -> PromptUpdateDetails: + types: torch.Tensor = out_mm_kwargs["image"][item_idx]["types"].data + full = (IMAGE_SENTINEL_BASE_ID + types).tolist() + return PromptUpdateDetails.select_token_id(full, image_embed_id) + + return [ + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_image_replacement, + ), + ] + + def _apply_token_matches_with_placeholders( + self, + token_ids: list[int], + mm_prompt_updates: MultiModalPromptUpdates, + ) -> tuple[ + list[int], + MultiModalPromptUpdatesApplyResult, + Mapping[str, list[PlaceholderFeaturesInfo]], + ]: + # Same as the base implementation, except that each image block gets + # its compressor-alignment pad (``3 - start_pos % 4`` IMAGE_PAD + # sentinels) prepended while splicing: the pad depends on the block's + # final position in the prompt, which is unknown when the (cacheable) + # prompt updates are built. + matched_updates, result = _plan_prompt_updates(token_ids, mm_prompt_updates) + placeholders: dict[str, list[PlaceholderFeaturesInfo]] = { + modality: [] for modality in mm_prompt_updates + } + + pad_id = IMAGE_SENTINEL_BASE_ID + IMAGE_PAD + + new_token_ids = list[int]() + prev_end_idx = 0 + for matched_update in matched_updates: + update = matched_update.update + match = matched_update.match + + if update.mode == UpdateMode.INSERT: + end_idx_to_insert = match.end_idx + elif update.mode == UpdateMode.REPLACE: + end_idx_to_insert = match.start_idx + else: + assert_never(update.mode) + + new_token_ids.extend(token_ids[prev_end_idx:end_idx_to_insert]) + start_idx = len(new_token_ids) + + tokens = list(update.content.full) + if tokens and update.modality == "image": + compress_pad = COMPRESS_PAD_TO - 1 - start_idx % COMPRESS_PAD_TO + tokens = [pad_id] * compress_pad + tokens + + if tokens: + content_is_embed = update.content.is_embed + is_embed = ( + content_is_embed(tokens) if content_is_embed is not None else None + ) + placeholders[update.modality].append( + PlaceholderFeaturesInfo( + modality=update.modality, + item_idx=update.item_idx, + start_idx=start_idx, + tokens=tokens, + is_embed=is_embed, + ) + ) + new_token_ids.extend(tokens) + + prev_end_idx = match.end_idx + + new_token_ids.extend(token_ids[prev_end_idx:]) + return new_token_ids, result, placeholders diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 8aa5e5d64b..6838ef928d 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -536,10 +536,16 @@ def combine_topk_swa_indices( M: int, N: int, out: tuple[torch.Tensor, torch.Tensor] | None = None, + left_visible: torch.Tensor | None = None, + right_visible: torch.Tensor | None = None, + max_image_tokens: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: num_tokens = topk_indices.shape[0] + # max_image_tokens widens the SWA column region for in-image + # bidirectional visibility; the width is fixed per model so the + # caller-provided workspace stays valid on image-free batches. combined_topk = ( - (topk + window_size + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) + (topk + window_size + max_image_tokens + _SPARSE_PREFILL_TOPK_ALIGNMENT - 1) // _SPARSE_PREFILL_TOPK_ALIGNMENT * _SPARSE_PREFILL_TOPK_ALIGNMENT ) @@ -568,6 +574,9 @@ def combine_topk_swa_indices( TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, + left_visible=left_visible, + right_visible=right_visible, + max_image_tokens=max_image_tokens, ) return combined_indices, combined_lens @@ -628,6 +637,7 @@ class CombineTopkSwaIndicesKernel( TOP_K: int COMPRESS_RATIO: int WINDOW_SIZE: int + IMAGE_WIDTH: int PADDED_TOP_K: int input_variant: TritonPointerInputVariant @@ -649,11 +659,14 @@ class CombineTopkSwaIndicesKernel( query_start_loc_ptr, seq_lens_ptr, gather_lens_ptr, + left_visible_ptr, + right_visible_ptr, M, N, TOP_K: tl.constexpr, COMPRESS_RATIO: tl.constexpr, WINDOW_SIZE: tl.constexpr, + IMAGE_WIDTH: tl.constexpr, PADDED_TOP_K: tl.constexpr, ): batch_idx = tl.program_id(0) @@ -682,7 +695,19 @@ class CombineTopkSwaIndicesKernel( token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) + if IMAGE_WIDTH > 0: + # In-image bidirectional visibility: the window starts up to + # max(left - (window - 1), 0) positions earlier and extends + # `right` positions past the query token (both 0 outside + # image spans, reducing to the plain causal window). + left = tl.load(left_visible_ptr + token_idx) + right = tl.load(right_visible_ptr + token_idx) + else: + left = 0 + right = 0 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) + swa_start = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) + swa_len = pos + right - swa_start + 1 offset = tl.arange(0, PADDED_TOP_K) mask = offset < topk_len @@ -695,18 +720,19 @@ class CombineTopkSwaIndicesKernel( topk_indices + M * batch_idx, mask=mask, ) - offset = tl.arange(0, WINDOW_SIZE) # Index into gathered buffer: N + (position - gather_start) - # For positions [pos - swa_len + 1, pos], the buffer indices are: - # [N + pos - swa_len + 1 - gather_start, N + pos - gather_start] - tl.store( - combined_indices_ptr - + token_idx * combined_indices_stride - + topk_len - + offset, - M * batch_idx + N + offset + pos - swa_len + 1 - gather_start, - mask=offset < swa_len, - ) + # For positions [swa_start, pos + right], the buffer indices are: + # [N + swa_start - gather_start, N + pos + right - gather_start] + for i in range(0, WINDOW_SIZE + IMAGE_WIDTH, WINDOW_SIZE): + swa_offset = i + tl.arange(0, WINDOW_SIZE) + tl.store( + combined_indices_ptr + + token_idx * combined_indices_stride + + topk_len + + swa_offset, + M * batch_idx + N + swa_start + swa_offset - gather_start, + mask=swa_offset < swa_len, + ) combined_len = topk_len + swa_len tl.store(combined_lens_ptr + token_idx, combined_len) @@ -722,6 +748,7 @@ class CombineTopkSwaIndicesKernel( topk: int, compress_ratio: int, WINDOW_SIZE: int, + image_width: int = 0, ) -> CompileKey: padded_topk = next_power_of_2(topk_width) input_variant = TritonPointerInputVariant.from_alignment( @@ -734,6 +761,7 @@ class CombineTopkSwaIndicesKernel( TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=WINDOW_SIZE, + IMAGE_WIDTH=image_width, PADDED_TOP_K=padded_topk, input_variant=input_variant, ) @@ -743,10 +771,19 @@ class CombineTopkSwaIndicesKernel( return [] window_size = _hf_config_int(vllm_config, "sliding_window", 128) + image_width = ( + _hf_config_int(vllm_config, "vision_max_n_token", 0) + if _hf_config_int(vllm_config, "vision_n_layers", 0) > 0 + else 0 + ) + # Warm both the plain-window variant (batches without image spans) + # and the in-image bidirectional variant. + image_widths = [0, image_width] if image_width > 0 else [0] return self._trace_dispatch(self.dispatch)( _DSV4_COMBINE_TOPK_SWA_WARMUP_INPUTS, _COMBINE_TOPK_SWA_POINTER_INPUTS, WINDOW_SIZE=window_size, + image_width=image_widths, ) def compile(self, compile_key: CompileKey) -> None: @@ -763,11 +800,14 @@ class CombineTopkSwaIndicesKernel( input_variant.pointer("query_start_loc", torch.int32), input_variant.pointer("seq_lens", torch.int32), input_variant.pointer("gather_lens", torch.int32), + int32_ptr, + int32_ptr, 1, # do not specialize M 1, # do not specialize N TOP_K=compile_key.TOP_K, COMPRESS_RATIO=compile_key.COMPRESS_RATIO, WINDOW_SIZE=compile_key.WINDOW_SIZE, + IMAGE_WIDTH=compile_key.IMAGE_WIDTH, PADDED_TOP_K=compile_key.PADDED_TOP_K, grid=(1, _COMBINE_TOPK_SWA_NUM_WORKERS), ) @@ -786,8 +826,13 @@ class CombineTopkSwaIndicesKernel( TOP_K: int, COMPRESS_RATIO: int, WINDOW_SIZE: int, + left_visible: torch.Tensor | None = None, + right_visible: torch.Tensor | None = None, + max_image_tokens: int = 0, ) -> None: num_reqs = seq_lens.shape[0] + has_image = left_visible is not None + image_width = max_image_tokens if has_image else 0 self.kernel[(num_reqs, _COMBINE_TOPK_SWA_NUM_WORKERS)]( combined_indices, combined_indices.stride(0), @@ -797,11 +842,14 @@ class CombineTopkSwaIndicesKernel( query_start_loc, seq_lens, gather_lens, + left_visible if has_image else topk_indices, + right_visible if has_image else topk_indices, M, N, TOP_K=TOP_K, COMPRESS_RATIO=COMPRESS_RATIO, WINDOW_SIZE=WINDOW_SIZE, + IMAGE_WIDTH=image_width, PADDED_TOP_K=next_power_of_2(topk_indices.shape[-1]), ) @@ -828,20 +876,33 @@ def build_flashinfer_mixed_sparse_indices( decode_is_valid_token: torch.Tensor | None = None, swa_block_span: int | None = None, compressed_block_span: int | None = None, + prefill_left_visible: torch.Tensor | None = None, + prefill_right_visible: torch.Tensor | None = None, + max_image_tokens: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. - Produces ``sparse_indices`` of shape ``[num_tokens, swa_index_width + - padded_topk]`` (the first ``swa_index_width`` columns are SWA slot ids, the + Produces ``sparse_indices`` of shape ``[num_tokens, swa_total_width + + padded_topk]`` (the first ``swa_total_width`` columns are SWA slot ids, the rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length per token). Decode tokens read precomputed SWA/compressed indices; prefill tokens derive their SWA window from the position and translate local compressed indices to global slots via the block tables. + + When ``prefill_left_visible``/``prefill_right_visible`` are given (vision + variant), the SWA column region widens by ``max_image_tokens`` and prefill + tokens inside an image span get a bidirectionally widened window; decode + rows are padded with -1 across the extra columns. """ assert decode_swa_indices.dtype == torch.int32 assert decode_swa_indices.dim() == 2 swa_index_width = decode_swa_indices.shape[-1] assert swa_index_width >= window_size + has_image = prefill_left_visible is not None + image_width = max_image_tokens if has_image else 0 + swa_total_width = swa_index_width + image_width + if has_image: + assert prefill_right_visible is not None if decode_compressed_topk_lens is not None: assert decode_compressed_topk_lens.dtype == torch.int32 assert prefill_topk_indices.dtype == torch.int32 @@ -890,7 +951,7 @@ def build_flashinfer_mixed_sparse_indices( padded_topk = max(topk, decode_compressed_topk) padded_topk = (padded_topk + 3) // 4 * 4 sparse_indices = torch.empty( - (num_tokens, swa_index_width + padded_topk), + (num_tokens, swa_total_width + padded_topk), dtype=torch.int32, device=decode_swa_indices.device, ) @@ -924,6 +985,8 @@ def build_flashinfer_mixed_sparse_indices( decode_is_valid_token, prefill_topk_indices, prefill_topk_indices.stride(0), + prefill_left_visible if has_image else token_to_req_indices, + prefill_right_visible if has_image else token_to_req_indices, query_start_loc, seq_lens, token_to_req_indices, @@ -938,6 +1001,7 @@ def build_flashinfer_mixed_sparse_indices( NUM_DECODE_TOKENS=num_decode_tokens, WINDOW_SIZE=window_size, SWA_INDEX_WIDTH=swa_index_width, + IMAGE_WIDTH=image_width, COMPRESS_RATIO=compress_ratio, TOP_K=topk, PADDED_TOP_K=padded_topk, @@ -992,6 +1056,8 @@ def _build_flashinfer_mixed_sparse_indices_kernel( decode_is_valid_token_ptr, prefill_topk_indices_ptr, prefill_topk_stride, + left_visible_ptr, + right_visible_ptr, query_start_loc_ptr, seq_lens_ptr, token_to_req_indices_ptr, @@ -1006,6 +1072,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( NUM_DECODE_TOKENS, WINDOW_SIZE: tl.constexpr, SWA_INDEX_WIDTH: tl.constexpr, + IMAGE_WIDTH: tl.constexpr, COMPRESS_RATIO: tl.constexpr, TOP_K: tl.constexpr, PADDED_TOP_K: tl.constexpr, @@ -1017,14 +1084,18 @@ def _build_flashinfer_mixed_sparse_indices_kernel( TOPK_BLOCK_SIZE: tl.constexpr, ): token_idx = tl.program_id(0) + # Total SWA column count; > SWA_INDEX_WIDTH only for in-image + # bidirectional visibility (vision variant), where the extra columns are + # -1-padded for decode rows. + SWA_TOTAL_WIDTH: tl.constexpr = SWA_INDEX_WIDTH + IMAGE_WIDTH if token_idx < NUM_DECODE_TOKENS: - for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < SWA_INDEX_WIDTH + mask = offset < SWA_TOTAL_WIDTH values = tl.load( decode_swa_indices_ptr + token_idx * decode_swa_stride + offset, - mask=mask, + mask=offset < SWA_INDEX_WIDTH, other=-1, ) values = _remap_flashinfer_index(values, swa_block_size, swa_block_span) @@ -1067,7 +1138,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + SWA_INDEX_WIDTH + + SWA_TOTAL_WIDTH + offset, values, mask=mask, @@ -1081,7 +1152,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( else: compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) - tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len) return prefill_idx = token_idx - NUM_DECODE_TOKENS @@ -1093,13 +1164,22 @@ def _build_flashinfer_mixed_sparse_indices_kernel( start_pos = seq_len - query_len token_idx_in_query = token_idx - query_start pos = start_pos + token_idx_in_query - swa_len = tl.minimum(pos + 1, WINDOW_SIZE) - swa_start_pos = pos - swa_len + 1 + if IMAGE_WIDTH > 0: + # In-image bidirectional visibility: window starts up to + # max(left - (window - 1), 0) earlier and extends `right` past pos. + left = tl.load(left_visible_ptr + token_idx) + right = tl.load(right_visible_ptr + token_idx) + else: + left = 0 + right = 0 + left_add = tl.maximum(left - (WINDOW_SIZE - 1), 0) + swa_start_pos = tl.maximum(pos - (WINDOW_SIZE - 1) - left_add, 0) + swa_len = pos + right - swa_start_pos + 1 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < SWA_INDEX_WIDTH + mask = offset < SWA_TOTAL_WIDTH pos_offset = swa_start_pos + offset block_indices = pos_offset // swa_block_size block_numbers = tl.load( @@ -1143,10 +1223,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + SWA_INDEX_WIDTH + + SWA_TOTAL_WIDTH + offset, slot_ids, mask=mask, ) - tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len) diff --git a/vllm/models/deepseek_v4/common/vision.py b/vllm/models/deepseek_v4/common/vision.py new file mode 100644 index 0000000000..a54c35e52c --- /dev/null +++ b/vllm/models/deepseek_v4/common/vision.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision tower (ViT + aligner) with TP-sharded linears. + +Ported from the official reference implementation +(deepseek-ai/DeepSeek-V4-Flash-Vision-Exp). Weight names match the HF +checkpoint so no renaming is needed at load time. Attention and MLP weights +are tensor-parallel sharded (replicated when the vision head count is not +divisible by TP size, or under ``--mm-encoder-tp-mode data``); the patch +embed and norms are replicated, so the residual stream is full-width on +every rank. +""" + +from functools import lru_cache + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, +) +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.models.vision import is_vit_use_data_parallel + + +@lru_cache(8) +def get_vision_cos_sin( + n_h: int, n_w: int, dim: int, theta: float +) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w) + wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w) + freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() + freqs = (freqs * inv_freq).flatten(1) + return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) + + +def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x1, x2 = x.float().chunk(2, dim=-1) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype) + + +class DeepseekV4RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +class DeepseekV4PatchEmbed(nn.Module): + def __init__(self, config): + super().__init__() + # Replicated: the residual stream is full-width on every rank. + self.proj = ReplicatedLinear( + 3 * config.vision_patch_size**2, + config.vision_dim, + bias=True, + quant_config=None, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out, _ = self.proj(x.flatten(1)) + return out + + +class DeepseekV4VisionAttention(nn.Module): + def __init__(self, config, prefix: str = ""): + super().__init__() + # Convention from minimax_m3/siglip: compute the TP/DP choice locally + # instead of threading a flag through every constructor. + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) + self.tp_size = ( + 1 if use_data_parallel else get_tensor_model_parallel_world_size() + ) + self.n_heads = config.vision_n_heads // self.tp_size + self.head_dim = config.vision_dim // config.vision_n_heads + self.hidden_size = config.vision_dim + self.wqkv = QKVParallelLinear( + config.vision_dim, + self.head_dim, + config.vision_n_heads, + bias=True, + quant_config=None, + prefix=f"{prefix}.wqkv", + disable_tp=use_data_parallel, + ) + self.wo = RowParallelLinear( + config.vision_dim, + config.vision_dim, + bias=True, + quant_config=None, + prefix=f"{prefix}.wo", + disable_tp=use_data_parallel, + ) + self.attn = MMEncoderAttention( + num_heads=self.n_heads, + head_size=self.head_dim, + prefix=f"{prefix}.attn", + ) + + def forward( + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + n = x.size(0) + qkv, _ = self.wqkv(x) + q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1)) + q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d) + k = apply_rotary(k, cos, sin).unsqueeze(0) + # One image per call: a dense batch, no varlen packing metadata. + o = self.attn(q, k, v.unsqueeze(0)) + out, _ = self.wo(o.reshape(n, -1)) + return out + + +class DeepseekV4VisionMLP(nn.Module): + def __init__(self, config, prefix: str = ""): + super().__init__() + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) + self.w1 = MergedColumnParallelLinear( + config.vision_dim, + [config.vision_inter_dim] * 2, + bias=False, + quant_config=None, + prefix=f"{prefix}.w1", + disable_tp=use_data_parallel, + ) + self.w2 = RowParallelLinear( + config.vision_inter_dim, + config.vision_dim, + bias=False, + quant_config=None, + prefix=f"{prefix}.w2", + disable_tp=use_data_parallel, + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.w1(x) + out, _ = self.w2(self.act_fn(gate_up)) + return out + + +class DeepseekV4VisionBlock(nn.Module): + def __init__(self, config, prefix: str = ""): + super().__init__() + self.norm1 = DeepseekV4RMSNorm(config.vision_dim) + self.attn = DeepseekV4VisionAttention(config, prefix=f"{prefix}.attn") + self.norm2 = DeepseekV4RMSNorm(config.vision_dim) + self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp") + + def forward( + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + x = x + self.attn(self.norm1(x), cos, sin) + return x + self.mlp(self.norm2(x)) + + +class DeepseekV4ViT(nn.Module): + """DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE.""" + + def __init__(self, config): + super().__init__() + self.rope_dim = config.vision_dim // config.vision_n_heads // 2 + self.rope_theta = config.vision_rope_theta + self.patch_embed = DeepseekV4PatchEmbed(config) + self.blocks = nn.ModuleList( + [ + DeepseekV4VisionBlock(config, prefix=f"blocks.{i}") + for i in range(config.vision_n_layers) + ] + ) + self.norm = DeepseekV4RMSNorm(config.vision_dim) + + def forward( + self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int + ) -> torch.Tensor: + x = self.patch_embed(patches) + cos, sin = get_vision_cos_sin(n_vit_h, n_vit_w, self.rope_dim, self.rope_theta) + cos = cos.to(device=x.device) + sin = sin.to(device=x.device) + for block in self.blocks: + x = block(x, cos, sin) + return self.norm(x) + + +class DeepseekV4Aligner(nn.Module): + """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.""" + + def __init__(self, config): + super().__init__() + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) + self.downsample_ratio = config.vision_downsample_ratio + in_dim = config.vision_dim * self.downsample_ratio**2 + self.w1 = ColumnParallelLinear( + in_dim, + config.hidden_size, + bias=True, + quant_config=None, + disable_tp=use_data_parallel, + ) + self.w2 = RowParallelLinear( + config.hidden_size, + config.hidden_size, + bias=True, + quant_config=None, + disable_tp=use_data_parallel, + ) + + def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: + r = self.downsample_ratio + x = x.view(n_vit_h, n_vit_w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -n_vit_w % r, 0, -n_vit_h % r)) + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1) + hidden, _ = self.w1(x) + out, _ = self.w2(F.gelu(hidden)) + return out diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 4ac362c95b..883ef5ecaf 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -435,6 +435,10 @@ class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention): decode_is_valid_token=decode_is_valid_token, swa_block_span=swa_block_span, compressed_block_span=compressed_block_span, + prefill_left_visible=swa_metadata.prefill_left_visible, + prefill_right_visible=swa_metadata.prefill_right_visible, + # getattr for tests that bypass __init__ via object.__new__. + max_image_tokens=getattr(self, "max_image_tokens", 0), ) if cache_key != "c4a": swa_metadata.flashinfer_sparse_index_cache[cache_key] = ( diff --git a/vllm/models/deepseek_v4/nvidia/flashmla.py b/vllm/models/deepseek_v4/nvidia/flashmla.py index b8eed8aefe..0ae0d87a14 100644 --- a/vllm/models/deepseek_v4/nvidia/flashmla.py +++ b/vllm/models/deepseek_v4/nvidia/flashmla.py @@ -119,7 +119,9 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): else: assert self.topk_indices_buffer is not None top_k = self.topk_indices_buffer.shape[-1] - combined_topk = round_up(top_k + self.window_size, 128) + combined_topk = round_up( + top_k + self.window_size + self.max_image_tokens, 128 + ) current_workspace_manager().get_simultaneous( ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), ((self.max_num_batched_tokens, combined_topk), torch.int32), @@ -311,7 +313,7 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): ) assert chunk_plan, "prefill chunk plan must be non-empty when num_prefills > 0" workspace_manager = current_workspace_manager() - combined_topk = round_up(top_k + self.window_size, 128) + combined_topk = round_up(top_k + self.window_size + self.max_image_tokens, 128) for chunk_start, chunk_end, chunk_N, chunk_M in chunk_plan: chunk_size = chunk_end - chunk_start workspace = workspace_manager.get_simultaneous( @@ -369,6 +371,21 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): chunk_M, chunk_N, out=(combined_indices_out, combined_lens_out), + left_visible=( + swa_metadata.prefill_left_visible[ + num_decode_tokens + query_start : num_decode_tokens + query_end + ] + if swa_metadata.prefill_left_visible is not None + else None + ), + right_visible=( + swa_metadata.prefill_right_visible[ + num_decode_tokens + query_start : num_decode_tokens + query_end + ] + if swa_metadata.prefill_right_visible is not None + else None + ), + max_image_tokens=self.max_image_tokens, ) flash_mla_sparse_fwd( q=q[query_start:query_end], diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index a469526316..3247877db0 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -92,6 +92,8 @@ from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID + logger = init_logger(__name__) @@ -798,6 +800,12 @@ class DeepseekV4MoE(nn.Module): self.gate.e_score_correction_bias = None self.gate.tid2eid = None + self.gate.bias_vl = None + # Image tokens borrow five consecutive reserved in-vocab ids starting + # at IMAGE_SENTINEL_BASE_ID; 0 disables vision routing (text model). + self.image_sentinel_lo = ( + IMAGE_SENTINEL_BASE_ID if getattr(config, "vision_n_layers", 0) > 0 else 0 + ) is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 if is_hash_moe: @@ -813,12 +821,25 @@ class DeepseekV4MoE(nn.Module): ), requires_grad=False, ) - elif getattr(config, "topk_method", None) == "noaux_tc": + if getattr(config, "topk_method", None) == "noaux_tc" and ( + not is_hash_moe or getattr(config, "vision_n_layers", 0) > 0 + ): + # Vision checkpoints ship a gate bias on hash layers too (it is + # unused for routing there; image tokens use bias_vl instead). self.gate.e_score_correction_bias = nn.Parameter( torch.empty(config.n_routed_experts, dtype=torch.float32), requires_grad=False, ) + if getattr(config, "vision_n_layers", 0) > 0: + # Vision checkpoints route image sentinel tokens with bias_vl + # instead of e_score_correction_bias / the hash table. Created on + # every MoE layer, hash layers included. + self.gate.bias_vl = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + if config.n_shared_experts is None: self.shared_experts = None else: @@ -950,6 +971,8 @@ class DeepseekV4MoE(nn.Module): routed_scaling_factor=self.routed_scaling_factor, e_score_correction_bias=self.gate.e_score_correction_bias, hash_indices_table=self.gate.tid2eid, + bias_vl=getattr(self.gate, "bias_vl", None), + image_sentinel_lo=self.image_sentinel_lo, swiglu_limit=self.swiglu_limit, router_logits_dtype=torch.float32, enable_eplb=parallel_config.enable_eplb, @@ -962,6 +985,10 @@ class DeepseekV4MoE(nn.Module): ) -> torch.Tensor: if self.gate.tid2eid is not None and input_ids is None: raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + # getattr for test doubles that fake the gate module. + bias_vl = getattr(self.gate, "bias_vl", None) + if bias_vl is not None and input_ids is None: + raise ValueError("DeepSeek V4 vision MoE routing requires input_ids.") if not self.use_mega_moe: return self._forward_fused_moe(hidden_states, input_ids) @@ -981,6 +1008,8 @@ class DeepseekV4MoE(nn.Module): input_tokens=input_ids, hash_indices_table=self.gate.tid2eid, routed_scaling_factor=self.routed_scaling_factor, + bias_vl=bias_vl.data if bias_vl is not None else None, + image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, ) activation_clamp = ( float(self.swiglu_limit) if self.swiglu_limit is not None else None diff --git a/vllm/models/deepseek_v4/nvidia/mtp.py b/vllm/models/deepseek_v4/nvidia/mtp.py index 7dd32fe972..b35ec6481a 100644 --- a/vllm/models/deepseek_v4/nvidia/mtp.py +++ b/vllm/models/deepseek_v4/nvidia/mtp.py @@ -178,7 +178,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module): inputs_embeds ).unsqueeze(-2) hidden_states, residual, post_mix, res_mix = self.mtp_block( - positions=positions, x=hidden_states, input_ids=None + positions=positions, x=hidden_states, input_ids=input_ids ) hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) if self.mtp_block.use_sequence_parallel: diff --git a/vllm/models/deepseek_v4/nvidia/vl_model.py b/vllm/models/deepseek_v4/nvidia/vl_model.py new file mode 100644 index 0000000000..bd40b7bce9 --- /dev/null +++ b/vllm/models/deepseek_v4/nvidia/vl_model.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). + +Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: + +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE + sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / + ``image_newline`` / ``image_end``) fill the remaining sentinel positions. +- Image placeholders (``<|deepseek_image|>``) are expanded by the processor + in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved + in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` + (see ``common/vision.py`` for the tower itself). +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before + its hyper-connection stream expansion. Raw ``input_ids`` still flow into + the model so the MoE router can apply ``bias_vl`` to image tokens + (``requires_raw_input_tokens``). +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + +from ..common.mm_preprocess import ( + IMAGE_PLACEHOLDER, + IMAGE_SENTINEL_BASE_ID, + DeepseekV4VLDummyInputsBuilder, + DeepseekV4VLMultiModalProcessor, + DeepseekV4VLProcessingInfo, + image_sentinel_mask, +) +from ..common.vision import DeepseekV4Aligner, DeepseekV4ViT +from .model import _make_deepseek_v4_weights_mapper + + +def _make_deepseek_v4_vl_weights_mapper( + expert_dtype: str, image_enabled: bool +) -> WeightsMapper: + """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" + base = _make_deepseek_v4_weights_mapper(expert_dtype) + orig_to_new_prefix: dict[str, str | None] = { + "layers.": "language_model.model.layers.", + "embed.": "language_model.model.embed.", + "norm.": "language_model.model.norm.", + "hc_head": "language_model.model.hc_head", + "mtp.": "language_model.model.mtp.", + } + if not image_enabled: + orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) + return WeightsMapper( + orig_to_new_prefix=orig_to_new_prefix, + orig_to_new_regex=base.orig_to_new_regex, + orig_to_new_suffix={ + "head.weight": "language_model.lm_head.weight", + "embed.weight": "embed_tokens.weight", + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", + }, + orig_to_new_substr={ + ".shared_experts.w2": ".shared_experts.down_proj", + # The MTP/DSpark draft heads are not supported for the vision + # variant; drop their weights. + "mtp.": None, + }, + ) + + +@MULTIMODAL_REGISTRY.register_processor( + DeepseekV4VLMultiModalProcessor, + info=DeepseekV4VLProcessingInfo, + dummy_inputs=DeepseekV4VLDummyInputsBuilder, +) +class DeepseekV4ForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 +): + """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. + + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) + delegates through ``language_model`` via the protocol defaults. + """ + + # The MoE router needs raw token ids to detect image sentinel tokens + # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. + requires_raw_input_tokens = True + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return IMAGE_PLACEHOLDER + raise ValueError(f"Unsupported modality: {modality!r}") + + def __init__(self, *, vllm_config, prefix: str = "") -> None: + super().__init__() + model_config = vllm_config.model_config + config = model_config.hf_config + self.config = config + self.multimodal_config = model_config.multimodal_config + assert self.multimodal_config is not None + + image_enabled = ( + config.vision_n_layers > 0 + and self.multimodal_config.get_limit_per_prompt("image") > 0 + ) + with self._mark_tower_model(vllm_config, {"image"}): + self.vision: DeepseekV4ViT | None = None + self.aligner: DeepseekV4Aligner | None = None + self.image_start: nn.Parameter | None = None + self.image_end: nn.Parameter | None = None + self.image_newline: nn.Parameter | None = None + self.image_pad: nn.Parameter | None = None + if image_enabled: + self.vision = DeepseekV4ViT(config) + self.aligner = DeepseekV4Aligner(config) + for name in ( + "image_start", + "image_end", + "image_newline", + "image_pad", + ): + setattr( + self, + name, + nn.Parameter( + torch.empty(config.hidden_size, dtype=torch.float32) + ), + ) + self.vision.to(dtype=model_config.dtype) + self.aligner.to(dtype=model_config.dtype) + + with self._mark_language_model(vllm_config): + # The arch convertor routes any config with a vision tower to + # this wrapper class; mark the copy handed to the text backbone + # so it resolves to DeepseekV4ForCausalLM instead of recursing + # (with_hf_config deepcopies the config, the marker survives). + config._dsv4_vl_inner = True # type: ignore[attr-defined] + try: + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["DeepseekV4ForCausalLM"], + ) + finally: + del config._dsv4_vl_inner # type: ignore[attr-defined] + # The outer mapper (see load_weights) fully resolves HF names into + # this wrapper's namespace before AutoWeightsLoader strips the + # "language_model." prefix and delegates to the child's load_weights, + # so the child's own mapper must be a no-op. Its suffix rules are not + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would + # re-fire "head.weight" -> "lm_head.weight"). + self.language_model.hf_to_vllm_mapper = WeightsMapper() + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + expert_dtype = getattr(config, "expert_dtype", "fp4") + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( + expert_dtype, image_enabled + ) + + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: + patches = kwargs.pop("patches", None) + if patches is None: + return None + vit_grid = kwargs.pop("vit_grid", None) + llm_grid = kwargs.pop("llm_grid", None) + perm = kwargs.pop("perm", None) + assert vit_grid is not None and llm_grid is not None and perm is not None + return { + "patches": patches, + "vit_grid": vit_grid, + "llm_grid": llm_grid, + "perm": perm, + } + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + perm: torch.Tensor, + ) -> torch.Tensor: + assert self.vision is not None and self.aligner is not None + image_embeds = self.aligner( + self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w + ) + # Reorder into the N-layout block order used in the prompt. + return image_embeds[perm.to(image_embeds.device)] + + def _process_image_input( + self, + patches: torch.Tensor, + vit_grid: torch.Tensor, + llm_grid: torch.Tensor, + perm: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + assert self.vision is not None and self.aligner is not None + patches = patches.to(self.aligner.w1.weight.dtype) + + embeds: list[torch.Tensor] = [] + vit_offset = 0 + llm_offset = 0 + for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( + vit_grid.tolist(), llm_grid.tolist(), strict=True + ): + n_vit = n_vit_h * n_vit_w + n_llm = n_llm_h * n_llm_w + embeds.append( + self._encode_image( + patches[vit_offset : vit_offset + n_vit], + n_vit_h, + n_vit_w, + perm[llm_offset : llm_offset + n_llm], + ) + ) + vit_offset += n_vit + llm_offset += n_llm + return tuple(embeds) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None or self.vision is None: + return [] + return self._process_image_input( + image_input["patches"], + image_input["vit_grid"], + image_input["llm_grid"], + image_input["perm"], + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # All ids are in-vocab here: image-block sentinels are borrowed + # reserved tokens (their embedding rows are always overwritten below). + inputs_embeds = self.language_model.embed_input_ids(input_ids) + + if self.image_start is not None: + # Branch-free sentinel overwrite: safe inside compiled/captured + # regions (no data-dependent control flow). + sentinel_mask = image_sentinel_mask(input_ids) + if is_multimodal is not None: + # IMAGE positions get vision embeddings via the merge below. + sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) + table = torch.stack( + [ + self.image_start, + self.image_pad, + self.image_pad, + self.image_newline, + self.image_end, + ] + ).to(inputs_embeds.dtype) + idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) + inputs_embeds = torch.where( + sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.compute_logits_local(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.language_model.get_expert_mapping() + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" + return self.language_model.get_mtp_target_hidden_states() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Map HF names into this wrapper's namespace up front and sort, so + # the "language_model." group reaches the child loader as one + # contiguous block (AutoWeightsLoader delegates per contiguous group, + # and the child's load_weights finalizes fused expert weights, which + # must not run on a partially loaded model). + mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0]) + loader = AutoWeightsLoader(self) + loaded_params = loader.load_weights(mapped) + # The child's load_weights already ran its post-load finalization. + self._weights_finalized = True + return loaded_params + + def process_weights_after_loading(self) -> None: + # Model-level post-load hook (called by the loader after any load + # format). Under DummyModelLoader the child's load_weights — and + # hence its finalize step — is bypassed, so run it here instead. + if getattr(self, "_weights_finalized", False): + return + self.language_model.process_weights_after_loading() diff --git a/vllm/models/deepseek_v4/vl_stub.py b/vllm/models/deepseek_v4/vl_stub.py new file mode 100644 index 0000000000..0f15a82f01 --- /dev/null +++ b/vllm/models/deepseek_v4/vl_stub.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stub for platforms where the DeepSeek-V4 vision variant is unsupported.""" + +from torch import nn + + +class DeepseekV4ForConditionalGeneration(nn.Module): + def __init__(self, *, vllm_config, prefix: str = ""): + super().__init__() + raise NotImplementedError( + "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is only " + "supported on NVIDIA GPUs for now." + ) diff --git a/vllm/tokenizers/deepseek_v4_encoding.py b/vllm/tokenizers/deepseek_v4_encoding.py index 1447c9de95..fb5bc44e85 100644 --- a/vllm/tokenizers/deepseek_v4_encoding.py +++ b/vllm/tokenizers/deepseek_v4_encoding.py @@ -28,6 +28,10 @@ USER_SP_TOKEN = "<|User|>" ASSISTANT_SP_TOKEN = "<|Assistant|>" LATEST_REMINDER_SP_TOKEN = "<|latest_reminder|>" +# Placeholder text inlined at each image's position; the multimodal processor +# later expands it into the sentinel token block. +IMAGE_PLACEHOLDER = "<|deepseek_image|>" + # Task special tokens for internal classification tasks DS_TASK_SP_TOKENS = { "action": "<|action|>", @@ -199,6 +203,30 @@ def find_last_user_index(messages: List[Dict[str, Any]]) -> int: # Message Rendering # ============================================================ +def flatten_content_blocks(content: Any) -> Any: + """Flatten OpenAI-style content blocks to plain text. + + Image blocks are inlined as IMAGE_PLACEHOLDER at their position. The image + data itself travels out of band (multi_modal_data) and is matched to + placeholders by order. Plain-string content passes through unchanged. + """ + if not isinstance(content, list): + return content + parts: List[str] = [] + for block in content: + if not isinstance(block, dict): + parts.append(str(block)) + elif block.get("type") in ("image", "image_url"): + parts.append(IMAGE_PLACEHOLDER) + elif block.get("type") == "text": + parts.append(block.get("text", "")) + else: + raise ValueError( + f"Unsupported content block type: {block.get('type')!r}" + ) + return "".join(parts) + + def render_message( index: int, messages: List[Dict[str, Any]], @@ -292,6 +320,8 @@ def render_message( block_type = block.get("type") if block_type == "text": parts.append(block.get("text", "")) + elif block_type in ("image", "image_url"): + parts.append(IMAGE_PLACEHOLDER) elif block_type == "tool_result": tool_content = block.get("content", "") if isinstance(tool_content, list): @@ -307,7 +337,7 @@ def render_message( parts.append(f"[Unsupported {block_type}]") prompt += "\n\n".join(parts) else: - prompt += content or "" + prompt += flatten_content_blocks(content) or "" elif role == "latest_reminder": 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]]: "content_blocks": [tool_block], }) elif role == "user": - text_block = {"type": "text", "text": msg.get("content", "")} + # Flatten OpenAI-style content blocks (images become placeholders) + text_content = flatten_content_blocks(msg.get("content", "")) + text_block = {"type": "text", "text": text_content} if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None: merged[-1]["content_blocks"].append(text_block) else: new_msg = { "role": "user", - "content": msg.get("content", ""), + "content": text_content, "content_blocks": [text_block], } # Preserve extra fields (task, wo_eos, mask, etc.) diff --git a/vllm/transformers_utils/configs/deepseek_v4.py b/vllm/transformers_utils/configs/deepseek_v4.py index 7708272c3b..15346e35e5 100755 --- a/vllm/transformers_utils/configs/deepseek_v4.py +++ b/vllm/transformers_utils/configs/deepseek_v4.py @@ -14,10 +14,39 @@ class DeepseekV4Config(PretrainedConfig): rope_scaling: dict[str, Any] | None = None, rope_parameters: dict[str, Any] | None = None, rope_theta: float = 10000.0, + vision_n_layers: int = 0, + vision_dim: int = 1024, + vision_n_heads: int = 16, + vision_inter_dim: int = 2816, + vision_patch_size: int = 14, + vision_rope_theta: float = 10000.0, + vision_downsample_ratio: int = 3, + vision_max_n_token: int = 384, + vision_min_pixels: int = 147456, + vision_max_wh_ratio: float = 8, **kwargs, ): self.max_position_embeddings = max_position_embeddings self.rope_scaling = rope_scaling self.rope_theta = rope_theta self.rope_parameters = rope_scaling or rope_parameters + # Vision tower config; vision_n_layers == 0 means text-only. + self.vision_n_layers = vision_n_layers + self.vision_dim = vision_dim + self.vision_n_heads = vision_n_heads + self.vision_inter_dim = vision_inter_dim + self.vision_patch_size = vision_patch_size + self.vision_rope_theta = vision_rope_theta + self.vision_downsample_ratio = vision_downsample_ratio + self.vision_max_n_token = vision_max_n_token + self.vision_min_pixels = vision_min_pixels + self.vision_max_wh_ratio = vision_max_wh_ratio + # The sparse-SWA index kernels widen the window within image spans + # in-kernel, so mm-prefix ranges longer than sliding_window must be + # kept (they are only consumed by those kernels). + self.mm_prefix_clamp_sliding_window = vision_n_layers > 0 + # The visibility span covers the sentinel block [IMAGE_START, + # IMAGE_END]; the mm placeholder additionally carries a leading + # compressor-alignment pad (see mm_preprocess.COMPRESS_PAD_TO). + self.mm_prefix_span_leading_pad_modulus = 4 if vision_n_layers > 0 else 0 super().__init__(**kwargs) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index d0e3d7c0dc..8baa9ee162 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -579,6 +579,40 @@ class DeepSeekMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) +class DeepseekV4ModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__( + self, + hf_config: PretrainedConfig, + hf_text_config: PretrainedConfig, + revision: str | None = None, + ): + # DeepSeek-V4-Flash-Vision-Exp ships the same architectures/model_type + # as the text-only DeepSeek-V4-Flash; route to the VL wrapper class + # when the config carries a vision tower. Mutate (not just override + # get_architectures) because model-class resolution reads the raw + # hf_config.architectures (get_model_architecture). + # Only rewrite the stock text architecture: speculative-draft configs + # (DSparkDraftModel / DeepSeekV4MTPModel) keep their own classes, and + # the VL wrapper marks the config copy it hands to the inner text + # backbone with _dsv4_vl_inner so this rewrite does not recurse. + if ( + getattr(hf_config, "vision_n_layers", 0) > 0 + and getattr(hf_config, "architectures", None) == ["DeepseekV4ForCausalLM"] + and not getattr(hf_config, "_dsv4_vl_inner", False) + ): + hf_config.architectures = ["DeepseekV4ForConditionalGeneration"] + super().__init__(hf_config, hf_text_config, revision) + + def is_mm_prefix_lm(self, supports_multimodal: bool = True) -> bool: + # The vision variant needs the mm-prefix plumbing: it makes the + # scheduler prefill image spans atomically (disable_chunked_mm_input) + # and routes per-request image ranges to the sparse-SWA metadata + # builder, which widens the sliding window bidirectionally in-kernel. + if not supports_multimodal: + return False + return getattr(self.hf_config, "vision_n_layers", 0) > 0 + + class MimoMTPModelArchConfigConvertor(ModelArchConfigConvertorBase): def get_num_hidden_layers(self) -> int: return getattr(self.hf_text_config, "num_nextn_predict_layers", 0) @@ -775,6 +809,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = { "cohere_asr": CohereAsrModelArchConfigConvertor, "dbrx": DbrxModelArchConfigConvertor, "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor, + "deepseek_v4": DeepseekV4ModelArchConfigConvertor, "diffusion_gemma_text": Gemma4ModelArchConfigConvertor, "ernie_mtp": ErnieMTPModelArchConfigConvertor, "falcon": FalconModelArchConfigConvertor, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index b968bda9fd..ccdf47f52e 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -9,6 +9,7 @@ from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, + zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, @@ -167,9 +168,15 @@ class DeepseekSparseSWAMetadata: decode_swa_width: int = 0 # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill). prefill_swa_indices: torch.Tensor | None = ( - None # [num_prefill_tokens, 1, window_size] + None # [num_prefill_tokens, 1, prefill_index_width] ) prefill_swa_lens: torch.Tensor | None = None # [num_prefill_tokens] + # In-image bidirectional visibility (vision variant): per-token counts of + # extra visible tokens to the left/right inside an image span. Indexed by + # absolute (decode-first) token position; only prefill rows are written. + # None when the model is text-only or the batch has no image spans. + prefill_left_visible: torch.Tensor | None = None + prefill_right_visible: torch.Tensor | None = None # Number of decode/prefill requests/tokens (batch is reordered: decodes first) num_decodes: int = 0 @@ -416,6 +423,17 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): assert hasattr(hf_config, "sliding_window") self.window_size = hf_config.sliding_window + # Vision variant: image spans (up to vision_max_n_token tokens) are + # visible bidirectionally, so prefill index rows widen from + # window_size to window_size + max_image_tokens. Text-only models keep + # max_image_tokens == 0 and take the original code paths everywhere. + self.max_image_tokens = ( + getattr(hf_config, "vision_max_n_token", 0) + if getattr(hf_config, "vision_n_layers", 0) > 0 + else 0 + ) + self.prefill_index_width = self.window_size + self.max_image_tokens + # Detect which DeepseekV4 layer types this model uses so we only build a # FlashMLA tile-scheduler plan for types that will actually be called. # Models without compress_ratios (pure SWA) fall back to swaonly. @@ -447,7 +465,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): self.prefill_swa_indices = torch.zeros( max_tokens, 1, - self.window_size, + self.prefill_index_width, dtype=torch.int32, device=self.device, ) @@ -456,6 +474,27 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): dtype=torch.int32, device=self.device, ) + # In-image visibility side buffers (vision variant only). The span + # CSR is uploaded per step from mm_req_doc_ranges; span capacity is + # bounded by max_tokens since each span holds at least one token. + if self.max_image_tokens > 0: + self.left_visible = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) + self.right_visible = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) + self.span_indptr = torch.zeros( + self.vllm_config.scheduler_config.max_num_seqs + 1, + dtype=torch.int32, + device=self.device, + ) + self.span_starts = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) + self.span_ends = torch.zeros( + max_tokens, dtype=torch.int32, device=self.device + ) self.is_valid_token = torch.zeros( max_tokens, dtype=torch.bool, @@ -556,6 +595,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): decode_swa_indices, self.decode_swa_lens, self.window_size, + decode_swa_indices.shape[-1], + self.decode_swa_lens, # unused (HAS_IMAGE=False) + self.decode_swa_lens, # unused (HAS_IMAGE=False) query_start_loc, seq_lens, token_to_req_indices, @@ -566,16 +608,42 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): token_offset=0, ) + # Vision variant: per-token in-image visibility for prefill tokens. + # Decode tokens are always past the image spans (spans are prefilled + # atomically), so the decode path above never needs them. + prefill_left_visible: torch.Tensor | None = None + prefill_right_visible: torch.Tensor | None = None + mm_ranges = common_attn_metadata.mm_req_doc_ranges + if ( + self.max_image_tokens > 0 + and num_prefill_tokens > 0 + and mm_ranges + and any(mm_ranges.values()) + ): + prefill_left_visible, prefill_right_visible = self._build_image_visibility( + common_attn_metadata.num_reqs, + mm_ranges, + num_decode_tokens, + num_prefill_tokens, + seq_lens, + query_start_loc, + token_to_req_indices, + ) + # Prefill SWA indices live in paged coordinates. `token_offset` lets # the kernel read is_valid_token / token_to_req_indices at absolute # prefill positions while writing output starting at index 0. if num_prefill_tokens > 0: + has_image = prefill_left_visible is not None prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens] prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens] _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( prefill_swa_indices, prefill_swa_lens, self.window_size, + self.prefill_index_width, + prefill_left_visible if has_image else prefill_swa_lens, + prefill_right_visible if has_image else prefill_swa_lens, query_start_loc, seq_lens, token_to_req_indices, @@ -584,6 +652,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): self.block_size, num_tokens=num_prefill_tokens, token_offset=num_decode_tokens, + has_image=has_image, ) # Pre-compute DeepseekV4 prefill metadata shared across all attention layers. @@ -623,6 +692,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): if num_prefill_tokens > 0 else None ), + prefill_left_visible=prefill_left_visible, + prefill_right_visible=prefill_right_visible, block_size=self.block_size, num_decodes=num_decodes, num_prefills=num_prefills, @@ -641,6 +712,50 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): **deepseek_v4_fields, # type: ignore[arg-type] ) + def _build_image_visibility( + self, + num_reqs: int, + mm_ranges: dict[int, list[tuple[int, int]]], + num_decode_tokens: int, + num_prefill_tokens: int, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + token_to_req_indices: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Upload image spans and compute per-token in-image visibility. + + Returns (left, right) int32 tensors indexed by absolute (decode-first) + token position; only prefill rows are meaningful. + """ + indptr = [0] * (num_reqs + 1) + starts: list[int] = [] + ends: list[int] = [] + for req_idx in range(num_reqs): + for span_start, span_end in mm_ranges.get(req_idx, ()): + starts.append(span_start) + ends.append(span_end) + indptr[req_idx + 1] = len(starts) + num_spans = len(starts) + assert num_spans <= self.span_starts.shape[0] + self.span_indptr[: num_reqs + 1].copy_(torch.tensor(indptr, dtype=torch.int32)) + if num_spans > 0: + self.span_starts[:num_spans].copy_(torch.tensor(starts, dtype=torch.int32)) + self.span_ends[:num_spans].copy_(torch.tensor(ends, dtype=torch.int32)) + _compute_image_visibility_kernel[(num_prefill_tokens,)]( + self.left_visible, + self.right_visible, + self.span_indptr, + self.span_starts, + self.span_ends, + query_start_loc, + seq_lens, + token_to_req_indices, + self.max_image_tokens, + token_offset=num_decode_tokens, + ) + end = num_decode_tokens + num_prefill_tokens + return self.left_visible[:end], self.right_visible[:end] + def update_draft_decode_metadata( self, metadata: DeepseekSparseSWAMetadata, @@ -658,6 +773,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): metadata.decode_swa_indices, metadata.decode_swa_lens, metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_indices.shape[-1], + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) + metadata.decode_swa_lens, # unused (HAS_IMAGE=False) metadata.query_start_loc, metadata.seq_lens, metadata.token_to_req_indices, @@ -755,90 +873,153 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): return result +@triton.jit(do_not_specialize=["token_offset"]) +def _compute_image_visibility_kernel( + left_visible_ptr, + right_visible_ptr, + span_indptr_ptr, + span_starts_ptr, + span_ends_ptr, + query_start_loc_ptr, + seq_lens_ptr, + token_to_req_indices_ptr, + max_image_tokens, + token_offset, +): + """Per-token in-image visible counts (port of `get_image_visible`). + + One program per prefill token. A token at position pos inside a span + [span_start, span_end] sees min(pos - span_start, max_image_tokens - 1) + extra tokens to its left and min(span_end - pos, max_image_tokens) to its + right; tokens outside every span get 0/0 (plain causal window). + """ + pid = tl.program_id(0) + token_idx = pid + token_offset + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + seq_len = tl.load(seq_lens_ptr + req_idx) + pos = seq_len - (query_end - query_start) + token_idx - query_start + + span_lo = tl.load(span_indptr_ptr + req_idx) + span_hi = tl.load(span_indptr_ptr + req_idx + 1) + left = tl.zeros((), dtype=tl.int32) + right = tl.zeros((), dtype=tl.int32) + for i in range(span_lo, span_hi): + span_start = tl.load(span_starts_ptr + i) + span_end = tl.load(span_ends_ptr + i) + in_span = (pos >= span_start) & (pos <= span_end) + left = tl.where( + in_span, tl.minimum(pos - span_start, max_image_tokens - 1), left + ) + right = tl.where(in_span, tl.minimum(span_end - pos, max_image_tokens), right) + tl.store(left_visible_ptr + token_idx, left) + tl.store(right_visible_ptr + token_idx, right) + + # TODO(ben): unify this kernel to reduce duplication +@triton.jit( + do_not_specialize=[ + "swa_indices_stride", + "block_table_stride", + "token_offset", + ] +) +def _compute_swa_indices_and_lens_kernel( + swa_indices_ptr, + swa_indices_stride, + swa_lens_ptr, + window_size, + index_width, + left_visible_ptr, + right_visible_ptr, + query_start_loc_ptr, + seq_lens_ptr, + token_to_req_indices_ptr, + is_valid_token_ptr, + block_table_ptr, + block_table_stride, + block_size, + token_offset, + HAS_IMAGE: tl.constexpr, + TRITON_BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + token_idx = pid + token_offset + is_valid = tl.load(is_valid_token_ptr + token_idx) + if not is_valid: + tl.store(swa_lens_ptr + pid, 0) + # Clear the row so a padded token cannot gather through stale indices. + for i in range(0, index_width, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + -1, + mask=offset < index_width, + ) + return + + req_idx = tl.load(token_to_req_indices_ptr + token_idx) + + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) + query_len = query_end - query_start + + seq_len = tl.load(seq_lens_ptr + req_idx) + prefix_len = seq_len - query_len + + pos = prefix_len + token_idx - query_start + if HAS_IMAGE: + # In-image bidirectional visibility widens the window: the window + # starts up to max(left - (window - 1), 0) positions earlier and + # extends `right` positions past the query token. + left = tl.load(left_visible_ptr + token_idx) + right = tl.load(right_visible_ptr + token_idx) + else: + left = 0 + right = 0 + left_add = tl.maximum(left - (window_size - 1), 0) + start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0) + end_pos = pos + right + 1 + + swa_len = end_pos - start_pos + tl.store(swa_lens_ptr + pid, swa_len) + + for i in range(0, index_width, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + + pos_offset = start_pos + offset + block_indices = pos_offset // block_size + block_numbers = tl.load( + block_table_ptr + req_idx * block_table_stride + block_indices, + mask=pos_offset < end_pos, + ) + block_offsets = pos_offset % block_size + slot_ids = block_numbers * block_size + block_offsets + + slot_ids = tl.where(offset < swa_len, slot_ids, -1) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + slot_ids, + mask=offset < index_width, + ) + + class ComputeSWAIndicesAndLensKernel( VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: window_size: int + index_width: int + has_image: int block_size: int triton_block_size: int - @staticmethod - @triton.jit( - do_not_specialize=[ - "swa_indices_stride", - "block_table_stride", - "token_offset", - ] - ) - def kernel( - swa_indices_ptr, - swa_indices_stride, - swa_lens_ptr, - window_size, - query_start_loc_ptr, - seq_lens_ptr, - token_to_req_indices_ptr, - is_valid_token_ptr, - block_table_ptr, - block_table_stride, - block_size, - token_offset, - TRITON_BLOCK_SIZE: tl.constexpr, - ): - pid = tl.program_id(0) - token_idx = pid + token_offset - is_valid = tl.load(is_valid_token_ptr + token_idx) - if not is_valid: - tl.store(swa_lens_ptr + pid, 0) - # Clear the row so a padded token cannot gather through stale indices. - for i in range(0, window_size, TRITON_BLOCK_SIZE): - offset = i + tl.arange(0, TRITON_BLOCK_SIZE) - tl.store( - swa_indices_ptr + pid * swa_indices_stride + offset, - -1, - mask=offset < window_size, - ) - return - - req_idx = tl.load(token_to_req_indices_ptr + token_idx) - - query_start = tl.load(query_start_loc_ptr + req_idx) - query_end = tl.load(query_start_loc_ptr + req_idx + 1) - query_len = query_end - query_start - - seq_len = tl.load(seq_lens_ptr + req_idx) - prefix_len = seq_len - query_len - - pos = prefix_len + token_idx - query_start - start_pos = tl.maximum(pos - window_size + 1, 0) - end_pos = pos + 1 - - swa_len = end_pos - start_pos - tl.store(swa_lens_ptr + pid, swa_len) - - for i in range(0, window_size, TRITON_BLOCK_SIZE): - offset = i + tl.arange(0, TRITON_BLOCK_SIZE) - - pos_offset = start_pos + offset - block_indices = pos_offset // block_size - block_numbers = tl.load( - block_table_ptr + req_idx * block_table_stride + block_indices, - mask=pos_offset < end_pos, - ) - block_offsets = pos_offset % block_size - slot_ids = block_numbers * block_size + block_offsets - - slot_ids = tl.where(offset < swa_len, slot_ids, -1) - tl.store( - swa_indices_ptr + pid * swa_indices_stride + offset, - slot_ids, - mask=offset < window_size, - ) + kernel = staticmethod(_compute_swa_indices_and_lens_kernel) def dispatch( # type: ignore[override] self, @@ -854,8 +1035,24 @@ class ComputeSWAIndicesAndLensKernel( *, window_size: int, block_size: int, + max_image_tokens: int = 0, ) -> list[CompileKey]: + # Decode rows always use the plain causal window. Vision models widen + # prefill rows by max_image_tokens and additionally need the in-image + # bidirectional variant for batches that contain image spans. + width_has_image = [(window_size, 0)] + if max_image_tokens > 0: + width_has_image += [ + (window_size + max_image_tokens, 0), + (window_size + max_image_tokens, 1), + ] return self._trace_dispatch(self.dispatch)( + zip_inputs( + *( + dict(index_width=index_width, has_image=has_image) + for index_width, has_image in width_has_image + ) + ), window_size=window_size, block_size=block_size, ) @@ -866,6 +1063,9 @@ class ComputeSWAIndicesAndLensKernel( swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), swa_lens=int32_ptr, window_size=compile_key.window_size, + index_width=compile_key.index_width, + left_visible=int32_ptr, + right_visible=int32_ptr, query_start_loc=int32_ptr, seq_lens=int32_ptr, token_to_req_indices=int32_ptr, @@ -874,6 +1074,7 @@ class ComputeSWAIndicesAndLensKernel( block_size=compile_key.block_size, num_tokens=1, token_offset=0, + has_image=bool(compile_key.has_image), ) @kernel_launcher @@ -882,6 +1083,9 @@ class ComputeSWAIndicesAndLensKernel( swa_indices: torch.Tensor, swa_lens: torch.Tensor, window_size: int, + index_width: int, + left_visible: torch.Tensor, + right_visible: torch.Tensor, query_start_loc: torch.Tensor, seq_lens: torch.Tensor, token_to_req_indices: torch.Tensor, @@ -891,10 +1095,12 @@ class ComputeSWAIndicesAndLensKernel( *, num_tokens: int, token_offset: int, + has_image: bool = False, ) -> LaunchSpec: return (num_tokens,), dict( swa_indices_stride=swa_indices.stride(0), block_table_stride=block_table.stride(0), + HAS_IMAGE=has_image, TRITON_BLOCK_SIZE=1024, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1f79e6ef44..d2271cc793 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2429,6 +2429,16 @@ class GPUModelRunner( _bidi_sw = getattr(hf_text_config, "sliding_window", None) _clamps_in_kernel = getattr( self.model, "mm_prefix_clamp_sliding_window", False + ) or getattr(hf_text_config, "mm_prefix_clamp_sliding_window", False) + # Some models (DeepSeek-V4 vision) define the bidirectional span + # over the whole sentinel block ([IMAGE_START, IMAGE_END]) rather + # than the embed tokens, and prepend a position-dependent + # alignment pad before the first sentinel. For those, derive the + # span from the full placeholder range and strip the pad. + # TODO(Isotr0py): Refactor mm_prefix_lm implementation + # for better readability and maintainability. + _span_pad_modulus = getattr( + hf_text_config, "mm_prefix_span_leading_pad_modulus", 0 ) for req_id in self.input_batch.req_ids: image_doc_ranges = [] @@ -2437,7 +2447,18 @@ class GPUModelRunner( if mm_feature.modality == "audio": continue pos_info = mm_feature.mm_position - img_doc_range = pos_info.extract_embeds_range() + if _span_pad_modulus: + pad = ( + _span_pad_modulus - 1 - pos_info.offset % _span_pad_modulus + ) + img_doc_range = [ + ( + pos_info.offset + pad, + pos_info.offset + pos_info.length - 1, + ) + ] + else: + img_doc_range = pos_info.extract_embeds_range() for r in img_doc_range: if ( not _clamps_in_kernel