iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Isotr0py
2026-09-02 11:09:09 -07:00
committed by GitHub
parent 488e6fd53c
commit 1356635d83
37 changed files with 2918 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/moe/moe_ops.h
+3 -1
View File
@@ -28,7 +28,9 @@ void topk_softplus_sqrt(
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding);
const std::optional<torch::stable::Tensor>& is_padding,
const std::optional<torch::stable::Tensor>& bias_vl,
int64_t image_sentinel_lo);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+172 -28
View File
@@ -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<int>(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<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
@@ -141,7 +215,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
config.numAttrs = 1;
const float scale = static_cast<float>(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<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -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<int64_t>(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<<<num_blocks, block_dim, 0, stream>>>(
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<NUM_EXPERTS, WARPS_PER_TB, 32, \
MAX_BYTES>( \
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<InputType, float>) {
if (use_hash && topk == 6 && renormalize &&
(num_experts == 256 || num_experts == 384)) {
launchDsv4HashTopk<IndType, HashIndType>(
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<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
const std::optional<torch::stable::Tensor>& is_padding) {
const std::optional<torch::stable::Tensor>& is_padding,
const std::optional<torch::stable::Tensor>& 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<float>();
}
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<float>();
}
auto launch = [&](auto* topk_indices_ptr) {
using OutIndType =
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
@@ -800,7 +912,8 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
tid2eid.value().const_data_ptr<int64_t>(), 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<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
tid2eid.value().const_data_ptr<int>(), 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<OutIndType, ComputeType,
int64_t>(
gating_output, topk_weights.mutable_data_ptr<float>(),
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, input_ids.value().const_data_ptr<int64_t>(),
static_cast<const int64_t*>(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<OutIndType, ComputeType,
int>(
gating_output, topk_weights.mutable_data_ptr<float>(),
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, input_ids.value().const_data_ptr<int>(),
static_cast<const int*>(nullptr), stream, is_padding_ptr,
bias_vl_ptr, image_sentinel_lo);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -818,7 +959,8 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr,
nullptr, 0);
}
};
@@ -843,7 +985,9 @@ void topk_softplus_sqrt(
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding) {
const std::optional<torch::stable::Tensor>& is_padding,
const std::optional<torch::stable::Tensor>& 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<float>(), 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<const __half*>(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<const __nv_bfloat16*>(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());
csrc/libtorch_stable/moe/torch_bindings.cpp
+2 -1
View File
@@ -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
docs/models/supported_models.md
+1
View File
@@ -528,6 +528,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
| `DeepseekVLV2ForCausalLM` | DeepSeek-VL2 | T + I<sup>+</sup> | `deepseek-ai/deepseek-vl2-tiny`, `deepseek-ai/deepseek-vl2-small`, `deepseek-ai/deepseek-vl2`, etc. | | ✅︎ |
| `DeepseekOCRForCausalLM` | DeepSeek-OCR | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR`, etc. | ✅︎ | ✅︎ |
| `DeepseekOCR2ForCausalLM` | DeepSeek-OCR-2 | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-OCR-2`, etc. | ✅︎ | ✅︎ |
| `DeepseekV4ForConditionalGeneration` | DeepSeek-V4 (vision) | T + I<sup>+</sup> | `deepseek-ai/DeepSeek-V4-Flash-Vision-Exp` | | ✅︎ |
| `Eagle2_5_VLForConditionalGeneration` | Eagle2.5-VL | T + I<sup>E+</sup> | `nvidia/Eagle2.5-8B`, etc. | ✅︎ | ✅︎ |
| `Ernie4_5_VLMoeForConditionalGeneration` | Ernie4.5-VL | T + I<sup>+</sup>/ V<sup>+</sup> | `baidu/ERNIE-4.5-VL-28B-A3B-PT`, `baidu/ERNIE-4.5-VL-424B-A47B-PT` | | ✅︎ |
| `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + I<sup>E+</sup> | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ |
tests/config/test_model_arch_config.py
+31
View File
@@ -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]
tests/kernels/moe/test_topk_softplus_sqrt.py
+257 -5
View File
@@ -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)
tests/models/multimodal/processing/test_tensor_schema.py
+5
View File
@@ -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(
tests/models/registry.py
+3
View File
@@ -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,
tests/models/test_initialization.py
+6
View File
@@ -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,
tests/models/test_registry.py
+5
View File
@@ -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
tests/tokenizers_/test_deepseek_v4.py
+58
View File
@@ -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]
tests/v1/attention/test_deepseek_v4_swa_visible.py
+580
View File
@@ -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
vllm/_custom_ops.py
+4
View File
@@ -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,
)
vllm/config/model.py
+4 -1
View File
@@ -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"
vllm/config/vllm.py
+1
View File
@@ -77,6 +77,7 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
"DeepseekV32MTPModel",
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
"DeepseekV4ForConditionalGeneration",
"DeepSeekV4MTPModel",
"Dots3NoteForCausalLM",
"Dots3NoteMTPModel",
vllm/model_executor/layers/fused_moe/layer.py
+7
View File
@@ -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:
vllm/model_executor/layers/fused_moe/router/dsv4_topk.py
+30
View File
@@ -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(),
)
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+68 -7
View File
@@ -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:
vllm/model_executor/layers/fused_moe/router/router_factory.py
+5
View File
@@ -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 (
vllm/model_executor/models/config.py
+1
View File
@@ -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
vllm/model_executor/models/registry.py
+4
View File
@@ -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",
vllm/models/deepseek_v4/__init__.py
+10
View File
@@ -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",
]
vllm/models/deepseek_v4/attention.py
+8
View File
@@ -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:
vllm/models/deepseek_v4/common/mm_preprocess.py
+512
View File
@@ -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
vllm/models/deepseek_v4/common/ops/cache_utils.py
+107 -27
View File
@@ -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)
vllm/models/deepseek_v4/common/vision.py
+232
View File
@@ -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
vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py
+4
View File
@@ -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] = (
vllm/models/deepseek_v4/nvidia/flashmla.py
+19 -2
View File
@@ -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],
vllm/models/deepseek_v4/nvidia/model.py
+30 -1
View File
@@ -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
vllm/models/deepseek_v4/nvidia/mtp.py
+1 -1
View File
@@ -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:
vllm/models/deepseek_v4/nvidia/vl_model.py
+333
View File
@@ -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()
vllm/models/deepseek_v4/vl_stub.py
+14
View File
@@ -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."
)
vllm/tokenizers/deepseek_v4_encoding.py
+35 -3
View File
@@ -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.)
vllm/transformers_utils/configs/deepseek_v4.py
+29
View File
@@ -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)
vllm/transformers_utils/model_arch_config_convertor.py
+35
View File
@@ -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,
vllm/v1/attention/backends/mla/sparse_swa.py
+280 -74
View File
@@ -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,
)
vllm/v1/worker/gpu_model_runner.py
+22 -1
View File
@@ -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