From 1c3e21b7328bb9894b62bfc4a3e8812611b7f4a8 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:02:56 +0000 Subject: [PATCH 01/15] deepseek_v4: DSpark speculative decoding under pipeline parallelism (SM80) Unlocks DSpark speculative decoding under pipeline parallelism on top of base commit f8ea5bb163 ("[Attention] DeepSeek-V4 sparse MLA on SM8x"). 8-file source diff baked into the published club-170hx image ghcr.io/pixelml/club-170hx:vllm-deepseek-v4-sm80-20260902, recovered from the build image (docker diff against a clean f8ea5bb16 checkout) and receipted at dsv4-vision-exp-4card-20260902T0440Z/receipts/vllm-source-patches.diff in the PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX evidence trail. This is the first commit that makes sm80 the exact source of the published text-path image, per the follow-up flagged in docs/SM80.md. CPU checks: py_compile and pyflakes clean on all 7 touched Python files. --- .dockerignore | 38 +--- vllm/config/speculative.py | 8 + .../layers/sparse_attn_indexer.py | 171 ++++++++++++++++-- vllm/models/deepseek_v4/nvidia/dspark.py | 10 + vllm/models/deepseek_v4/nvidia/model.py | 9 +- vllm/v1/worker/gpu/model_runner.py | 27 ++- vllm/v1/worker/gpu/pp_utils.py | 63 +++++++ .../v1/worker/gpu/spec_decode/dspark/utils.py | 90 ++++++++- 8 files changed, 360 insertions(+), 56 deletions(-) diff --git a/.dockerignore b/.dockerignore index 66447272e95a..ce6de41a2cd3 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,35 +1,3 @@ -/.venv -/build -dist -vllm/*.so -vllm/vllm-rs - -# Byte-compiled / optimized / DLL files -__pycache__/ -*.py[cod] -*$py.class - -.mypy_cache - -# Distribution / packaging -.Python -/build/ -cmake-build-*/ -CMakeUserPresets.json -develop-eggs/ -/dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -share/python-wheels/ -*.egg-info/ -.installed.cfg -*.egg -MANIFEST -rust/target/ +.git +__pycache__ +*.pyc diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 07bdcac68b84..568fa589812f 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1078,6 +1078,14 @@ def __post_init__(self): self.target_parallel_config, self.draft_tensor_parallel_size ) ) + if self.method == "dspark": + # The DSpark draft is NOT pipelined: the model runner builds + # it on the last PP rank only and it runs there whole, so it + # is a pp_size=1 model regardless of the target's split. + # Inheriting the target's pipeline_parallel_size would make + # verify_with_parallel_config demand SupportsPP from the + # draft architecture, which it neither implements nor needs. + self.draft_parallel_config.pipeline_parallel_size = 1 return self def _validate_suffix_decoding(self): diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index a0a9eeb7d5a6..4413d8097922 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Custom Sparse Attention Indexer layers.""" +import os + import torch import vllm.envs as envs @@ -75,6 +77,9 @@ def _assert_cutedsl_dcp_merge_supported( ) +_DSV4_LOGITS_ROW_CHUNK = int(os.environ.get("DSV4_LOGITS_ROW_CHUNK", "0")) + + def _merge_dcp_topk_global( logits: torch.Tensor, topk_indices: torch.Tensor, @@ -128,6 +133,78 @@ def _merge_dcp_topk_global( ) +def _top_k_per_row_prefill_torch( + logits: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + topk_indices: torch.Tensor, + topk_tokens: int, +) -> None: + """torch.topk fallback with ``top_k_per_row_prefill``'s output contract. + + Ported from vllm-project/vllm PR #49897, which added this for SM12x. The + CUDA kernel's histogram path (taken by rows with more than ``topk_tokens`` + candidates) can leave part of its dynamic-shared-memory output + uninitialized and copy it out as indices; downstream, + ``compute_global_topk_indices_and_lens`` treats any index ``>= 0`` as valid + and dereferences it into the KV block table, crashing with a CUDA illegal + memory access. + + Enabled here for SM8x as well: on 4x CMP 170HX this reproduces as an + ``Xid 31 MMU Fault ... ACCESS_TYPE_VIRT_WRITE`` that kills a worker on + prefills above roughly 128k tokens, while 123k passes. + + Select the top ``topk_tokens`` positions of ``logits[i, ks_i:ke_i)`` per + row, emit them relative to ``ks_i`` in ascending position order, and pad + rows with fewer candidates with ``-1`` (the kernel's short-row contract). + """ + num_cols = logits.shape[1] + ks = cu_seqlen_ks.to(torch.long)[:, None] + cols = torch.arange(num_cols, device=logits.device)[None, :] + valid = (cols >= ks) & (cols < cu_seqlen_ke.to(torch.long)[:, None]) + # logits is a per-chunk scratch buffer that is dead after top-k; mask it + # in place rather than materializing a masked copy. + logits.masked_fill_(~valid, float("-inf")) + k = min(topk_tokens, num_cols) + top_values, top_cols = logits.topk(k, dim=-1) + relative = (top_cols - ks).to(torch.int32) + # Downstream sparse-attention kernels iterate the selected KV positions in + # ascending order; emit position-sorted indices with the ``-1`` pads at the + # tail (torch.topk returns score order instead). Non-finite scores (rows + # shorter than ``topk_tokens``, or NaN logits) pad. + pad_sentinel = torch.iinfo(torch.int32).max + relative = torch.where( + top_values.isfinite(), relative, relative.new_full((), pad_sentinel) + ) + relative, _ = relative.sort(dim=-1) + relative = torch.where( + relative == pad_sentinel, relative.new_full((), -1), relative + ) + topk_indices[:, :k] = relative + if k < topk_tokens: + topk_indices[:, k:] = -1 + + +def _prefill_topk_needs_torch_fallback() -> bool: + """SM12x per PR #49897; SM8x added here (see _top_k_per_row_prefill_torch). + + Set VLLM_DSV4_PREFILL_TOPK_TORCH=0 to force the CUDA kernel back on, or =1 + to force the fallback on any architecture. + """ + import os + + forced = os.environ.get("VLLM_DSV4_PREFILL_TOPK_TORCH") + if forced is not None: + return forced == "1" + if not current_platform.is_cuda(): + return False + return current_platform.is_device_capability_family( + 120 + ) or current_platform.has_device_capability(80) and not ( + current_platform.has_device_capability(90) + ) + + @triton.jit def _fused_indexer_q_rope_quant_kernel( positions, @@ -515,6 +592,49 @@ def sparse_attn_indexer( cu_seqlen_ke, clean_logits=False, ) + elif _DSV4_LOGITS_ROW_CHUNK > 0 and dcp_world_size <= 1: + # SM80 Triton fallback, row-chunked. The [M, N] float32 + # logits transient is the indexer's largest allocation and + # scales with context; on the last PP rank it is what + # exhausts the device at long context. Each row's top-k + # reads only its own [ks, ke), so blocking over rows is + # exact and caps the transient at rows x N x 4 bytes. + _step = _DSV4_LOGITS_ROW_CHUNK + _tot = q_slice_cast.shape[0] + _torch_topk = _prefill_topk_needs_torch_fallback() + for _r0 in range(0, _tot, _step): + _r1 = min(_r0 + _step, _tot) + logits = fp8_mqa_logits_triton( + q_slice_cast[_r0:_r1], + (k_quant_cast, k_scale_cast), + weights[ + chunk.token_start + _r0 : chunk.token_start + _r1 + ], + cu_seqlen_ks[_r0:_r1], + cu_seqlen_ke[_r0:_r1], + clean_logits=False, + ) + if _torch_topk: + _top_k_per_row_prefill_torch( + logits, + cu_seqlen_ks[_r0:_r1], + cu_seqlen_ke[_r0:_r1], + topk_indices[_r0:_r1], + topk_tokens, + ) + else: + ops.top_k_per_row_prefill( + logits, + cu_seqlen_ks[_r0:_r1], + cu_seqlen_ke[_r0:_r1], + topk_indices[_r0:_r1], + _r1 - _r0, + logits.stride(0), + logits.stride(1), + topk_tokens, + ) + del logits + logits = q_slice_cast.new_empty((0, 0), dtype=torch.float32) else: # SM80/SM121 Triton fallback (DeepGEMM unavailable). logits = fp8_mqa_logits_triton( @@ -525,17 +645,28 @@ def sparse_attn_indexer( cu_seqlen_ke, clean_logits=False, ) - num_rows = logits.shape[0] - ops.top_k_per_row_prefill( - logits, - cu_seqlen_ks, - cu_seqlen_ke, - topk_indices, - num_rows, - logits.stride(0), - logits.stride(1), - topk_tokens, - ) + if logits.shape[0] == 0: + pass + elif _prefill_topk_needs_torch_fallback(): + num_rows = logits.shape[0] + _top_k_per_row_prefill_torch( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + topk_indices, + topk_tokens, + ) + else: + ops.top_k_per_row_prefill( + logits, + cu_seqlen_ks, + cu_seqlen_ke, + topk_indices, + logits.shape[0], + logits.stride(0), + logits.stride(1), + topk_tokens, + ) _merge_dcp_topk_global( logits, @@ -655,10 +786,20 @@ def sparse_attn_indexer( and current_platform.has_device_capability(90) and not current_platform.is_device_capability_family(120) ) - use_persistent_topk = current_platform.is_cuda() and topk_tokens in ( - 512, - 1024, - 2048, + # NOTE(sm8x): the radix `persistent_topk` returns WRONG indices when the + # candidate count lands between k and 2k. With index_topk=512 and + # compress_ratio=4 that is prompt length 2049-4096 -- an ordinary range -- + # and it does not crash, it emits fluent-looking degenerate text. The + # sibling `use_cooperative_topk` above carries a has_device_capability(90) + # gate; this one is missing it, so SM8x selects the broken kernel. Mirror + # the gate; SM8x then falls back to ops.top_k_per_row_decode. + # Reported on vllm-project/vllm#50576 by another CMP 170HX owner + # (boundary predicted 8/8; cudagraphs, chunked prefill and prompt + # repetition all ruled out); no throughput cost (19.5 -> 19.6 tok/s). + use_persistent_topk = ( + current_platform.is_cuda() + and topk_tokens in (512, 1024, 2048) + and current_platform.has_device_capability(90) ) if use_cooperative_topk: workspace_manager = current_workspace_manager() diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index bde088edae1a..adad56a7d617 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -452,7 +452,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict[name][: narrow.shape[0]].copy_(narrow) loaded_params.add(name) continue + # Vision-Exp checkpoint carries vision-language variants and + # hash-layer gate biases that the text-only DSpark draft does + # not define as parameters. Skip them the same way the target + # model loader does. + if name.endswith("bias_vl"): + continue if name.endswith(".ffn.gate.bias"): + import re as _re + _m = _re.search(r"layers\.(\d+)\.", name) + if _m and int(_m.group(1)) < self.config.num_hash_layers: + continue name = name.replace( ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" ) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 8806b66f97aa..db460dc8d270 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1216,6 +1216,13 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ) for name, loaded_weight in weights: + # Vision-Exp checkpoints store ffn.gate.bias on hash-MoE layers + # too, where routing is vocabulary-hash based and the fork + # intentionally registers no e_score_correction_bias parameter. + if name.endswith("ffn.gate.e_score_correction_bias"): + m = re.search(r"layers\.(\d+)\.", name) + if m and int(m.group(1)) < self.config.num_hash_layers: + continue if pad_shared_expert and ".shared_experts." in name: loaded_weight = self._pad_shared_expert_weight( self.quant_config, name, loaded_weight @@ -1507,7 +1514,7 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: return getattr(self.model, "_mtp_hidden_buffer", None) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self, skip_substrs=["mtp."]) + loader = AutoWeightsLoader(self, skip_substrs=["mtp.", "vision.", "aligner.", "image_", "_vl"]) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) def process_weights_after_loading(self) -> None: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 79dcc2cf20bd..0f35c32afed7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -209,7 +209,14 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): if self.speculative_config.method in ("eagle3", "dflash", "dspark"): # Drafting may require auxiliary hidden states from target model outputs self.use_aux_hidden_state_outputs = True - if self.use_pp: + # DSpark + PP is enabled here (see pp_utils.broadcast_draft and + # dspark/utils.load_dspark_model). It works because the drafter is + # already built on the last PP rank only (above), and for + # DeepSeek-V4 the aux-hidden-state taps (dspark_target_layer_ids + # = [40,41,42] of 43 layers) plus lm_head all land on that same + # last rank. eagle3/dflash keep the guard -- untested, and their + # aux layers are spread across ranks. + if self.use_pp and self.speculative_config.method != "dspark": raise ValueError( f"{self.speculative_config.method} with pipeline parallel " "is not supported." @@ -830,7 +837,19 @@ def update_pp_decode_requests(self): if self.pp_handler is not None: outputs = self.pp_handler.get_prev_sampled_outputs() if outputs is not None: + # Spec decode: scatter the relayed proposed draft tokens into this + # rank's state so the next step's combine_sampled_and_draft_tokens + # reads real values instead of zero-init (which gives ~0 acceptance + # and corrupt output). Pop before postprocess_sampled, which does + # not accept this kwarg. idx_mapping is -1 for excluded/freed reqs. + draft_tokens = outputs.pop("draft_tokens", None) + idx_mapping = outputs["idx_mapping"] self.postprocess_sampled(**outputs) + if draft_tokens is not None: + valid = idx_mapping >= 0 + self.req_states.draft_tokens[idx_mapping[valid]] = draft_tokens[ + valid + ] def add_requests(self, scheduler_output: SchedulerOutput) -> None: for new_req_data in scheduler_output.scheduled_new_reqs: @@ -1578,6 +1597,12 @@ def sample_tokens( mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens + if self.pp_handler is not None: + # Relay the proposed draft tokens to the non-last PP ranks so + # their next-step combine_sampled_and_draft_tokens reads real + # values instead of zero-init (otherwise acceptance ~= 0 and the + # output is garbage). Must be issued after propose(). + self.pp_handler.broadcast_draft(draft_tokens, input_batch) if self.num_speculative_steps > 0: # Spec-decode and diffusion LLMs both use draft tokens but the latter does diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 9b1786cbe4cf..38b45e45f1e6 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -30,6 +30,10 @@ class PendingRecv: # Snapshot of slot generation counters at receive time, used to # detect requests aborted since then. gen_at_receive_np: np.ndarray # [num_reqs] + # Spec decode: proposed draft tokens relayed from the last rank, scattered + # into the non-last rank's req_states.draft_tokens on consume. None when spec + # is off (max_sample_len == 1). + draft_tokens: torch.Tensor | None = None # [num_reqs, max_sample_len - 1] def compute_need_sampled_mask(input_batch: InputBatch) -> np.ndarray | None: @@ -121,6 +125,7 @@ def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: num_sampled=slot.num_sampled, num_rejected=slot.num_rejected, idx_mapping=idx_mapping, + draft_tokens=slot.draft_tokens, ) def receive(self, input_batch: InputBatch) -> bool: @@ -149,12 +154,28 @@ def receive(self, input_batch: InputBatch) -> bool: torch.distributed.broadcast( combined, src=self.last_rank, group=self.broadcast_group ) + # Spec decode: 3rd broadcast on this group — the proposed draft tokens + # relayed by the last rank (matches broadcast_draft's order). NCCL + # matches by op order on the communicator; the deferred event covers it. + draft_tokens = None + if self.max_sample_len > 1: + draft_tokens = torch.empty( + num_reqs, + self.max_sample_len - 1, + dtype=torch.int64, + device=self.device, + ) + torch.distributed.broadcast( + draft_tokens, src=self.last_rank, group=self.broadcast_group + ) event = self.broadcast_stream.record_event() num_sampled, num_rejected = combined.unbind(dim=0) # Must record_stream since these were allocated on broadcast stream but # later used on the main stream. sampled_tokens.record_stream(self.main_stream) combined.record_stream(self.main_stream) + if draft_tokens is not None: + draft_tokens.record_stream(self.main_stream) self.queue[-1] = PendingRecv( event, sampled_tokens, @@ -164,6 +185,7 @@ def receive(self, input_batch: InputBatch) -> bool: input_batch.idx_mapping_np, need_sampled_mask, gen_at_receive_np, + draft_tokens, ) return bool(need_sampled_mask.all()) @@ -180,6 +202,22 @@ def broadcast( return assert sampled_token_ids.dtype == torch.int64 + # The receiver always posts a [num_reqs, max_sample_len] buffer, but under + # spec decode the sampler emits width 1 on any step with no draft tokens + # (prefill, first decode) and width (num_spec+1) only once rejection + # sampling has run. Pad to max_sample_len so the NCCL broadcast element + # count matches the receiver on EVERY step (an unpadded width-1 send vs a + # width-(num_spec+1) recv is a count mismatch that deadlocks the receiver). + # Trailing -1 columns are placeholders ignored by post_update, which + # advances each request by its per-request num_sampled valid count. + width = sampled_token_ids.shape[-1] + if width != self.max_sample_len: + assert width < self.max_sample_len + padded = sampled_token_ids.new_full( + (sampled_token_ids.shape[0], self.max_sample_len), -1 + ) + padded[:, :width] = sampled_token_ids + sampled_token_ids = padded if current_platform.is_xpu(): self.main_stream.synchronize() @@ -197,3 +235,28 @@ def broadcast( ) for tensor in (sampled_token_ids, num_sampled, num_rejected): tensor.record_stream(self.broadcast_stream) + + def broadcast_draft( + self, draft_tokens: torch.Tensor, input_batch: InputBatch + ) -> None: + """Relay the proposed draft tokens (spec decode) from the last rank to the + non-last ranks. Issued AFTER propose(), so it is the 3rd broadcast on this + group for the step (sampled, combined, draft) — matching the order in which + `receive` posts its recvs. Gated identically to `broadcast` so op counts + stay matched across ranks. Without this, non-last ranks verify against a + zero-init req_states.draft_tokens (near-zero acceptance, corrupt output).""" + assert self.is_last_rank + if self.max_sample_len <= 1: + return + if compute_need_sampled_mask(input_batch) is None: + # No request needs sampled outputs next step; `broadcast` skipped too, + # so skip here to keep the per-step broadcast count matched. + return + draft_tokens = draft_tokens.to(torch.int64).contiguous() + with torch.cuda.stream(self.broadcast_stream): + # wait_stream so the side-stream broadcast sees propose()'s output. + self.broadcast_stream.wait_stream(self.main_stream) + torch.distributed.broadcast( + draft_tokens, src=self.last_rank, group=self.broadcast_group + ) + draft_tokens.record_stream(self.broadcast_stream) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 13c7644315d5..daf6d399fce5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -1,16 +1,81 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import os + +import torch import torch.nn as nn from vllm.config import VllmConfig, replace from vllm.distributed.parallel_state import get_pp_group +from vllm.logger import init_logger from vllm.model_executor.model_loader import get_model from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( _should_share, get_target_lm_head, ) +logger = init_logger(__name__) + +# Checkpoint names the token embedding can appear under, most specific first. +_EMBED_KEYS = ( + "embed.weight", + "model.embed_tokens.weight", + "model.embed.weight", + "embed_tokens.weight", +) + + +def _has_real_weight(module) -> bool: + """True for a materialised layer; False for PPMissingLayer / None. + + Under pipeline parallelism vLLM replaces the layers a rank does not own with + PPMissingLayer, which has no ``weight``. Aliasing one of those into the draft + silently produces a no-op layer rather than an error, so check explicitly. + """ + return module is not None and getattr(module, "weight", None) is not None + + +def _load_embed_from_checkpoint(embed: nn.Module, model_path: str) -> None: + """Fill the draft's own token embedding straight from the checkpoint. + + Only needed under PP: the drafter runs on the LAST pipeline rank, but the + target's ``embed_tokens`` lives on the FIRST, so there is nothing local to + alias. Reading the one tensor off disk (~1 GB) avoids adding a cross-rank + collective to model load, which would have to be ordered against every other + rank's initialisation. + """ + from safetensors import safe_open + + index_path = os.path.join(model_path, "model.safetensors.index.json") + key = shard = None + if os.path.exists(index_path): + with open(index_path) as f: + weight_map = json.load(f)["weight_map"] + for cand in _EMBED_KEYS: + if cand in weight_map: + key, shard = cand, os.path.join(model_path, weight_map[cand]) + break + if key is None: + raise RuntimeError( + f"DSpark+PP: could not find a token-embedding tensor in {index_path}; " + f"looked for {_EMBED_KEYS}. The draft needs its own copy because the " + "target's embedding lives on PP rank 0." + ) + + with safe_open(shard, framework="pt") as f: + w = f.get_tensor(key) + with torch.no_grad(): + # VocabParallelEmbedding pads the vocab dimension, so copy into the + # leading rows rather than assigning the whole tensor. + embed.weight.data[: w.shape[0]].copy_( + w.to(dtype=embed.weight.dtype, device=embed.weight.device) + ) + logger.info( + "DSpark+PP: loaded draft token embedding %s from key %r", tuple(w.shape), key + ) + def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: speculative_config = vllm_config.speculative_config @@ -46,8 +111,13 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo vllm_config=draft_vllm_config, model_config=draft_model_config ) - if get_pp_group().world_size != 1: - raise NotImplementedError("DSpark does not support pipeline parallelism.") + # NOTE(pp): DSpark used to hard-refuse pipeline parallelism here. It works + # provided (a) the drafter is built on the last PP rank only -- the model + # runner already does exactly that -- and (b) the draft can still reach a + # token embedding and an lm_head. For DeepSeek-V4 the aux-hidden-state taps + # (dspark_target_layer_ids = [40,41,42] of 43 layers) and lm_head both land + # on that same last rank, so only the embedding must be sourced separately. + is_pp = get_pp_group().world_size != 1 target_language_model = ( target_model.get_language_model() @@ -59,20 +129,32 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo target_embed = getattr(target_inner, "embed_tokens", None) draft_embed = getattr(draft_inner, "embed_tokens", None) - if target_embed is not None and _should_share( + if _has_real_weight(target_embed) and _should_share( draft_model, "has_own_embed_tokens", draft_embed, target_embed ): if draft_embed is not None: del draft_inner.embed_tokens draft_inner.embed_tokens = target_embed + elif is_pp and draft_embed is not None: + # PP: the target's embedding is a PPMissingLayer on this rank. Keep the + # draft's own module and populate it from the checkpoint -- the draft's + # load_weights() only consumes mtp.* tensors, so it is still empty here. + _load_embed_from_checkpoint(draft_embed, draft_model_config.model) target_lm_head = get_target_lm_head(target_model, target_language_model) draft_lm_head = getattr(draft_model, "lm_head", None) - if target_lm_head is not None and _should_share( + if _has_real_weight(target_lm_head) and _should_share( draft_model, "has_own_lm_head", draft_lm_head, target_lm_head ): if draft_lm_head is not None: del draft_model.lm_head draft_model.lm_head = target_lm_head + elif is_pp and not _has_real_weight(target_lm_head): + # Should not happen: the drafter runs on the last rank, which owns + # lm_head. Fail loudly rather than silently drafting through a no-op. + raise RuntimeError( + "DSpark+PP: the target lm_head is not materialised on this rank. The " + "drafter is expected to run on the last pipeline rank." + ) return draft_model From 7c03547433acf0b2898f92a8473aafe6dcc159ee Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:03:48 +0000 Subject: [PATCH 02/15] deepseek_v4: vision-only files + registry wiring (Path 3 vision port, step 1/4) Ports vision-only files unmodified from the pinned DeepSeek-V4-Flash-Vision-Exp head onto this SM80 fork, plus registry/ __init__ wiring to expose DeepseekV4ForConditionalGeneration as a multimodal model: - vllm/models/deepseek_v4/vl_stub.py (non-CUDA platform stub) - vllm/models/deepseek_v4/nvidia/vl_model.py - vllm/models/deepseek_v4/common/mm_preprocess.py - vllm/models/deepseek_v4/common/vision.py Source: PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX, branch codex/path3-vision-on-sm80-fork, format-patch 0001 (commit 2cb04cc1). No adaptation needed for these files: their imports resolve against generic vLLM infra already present on this fork, and the one deepseek_v4-internal symbol they need (_make_deepseek_v4_weights_mapper) already exists in nvidia/model.py. CPU checks: py_compile and pyflakes clean on all 6 touched/added files. --- vllm/model_executor/models/registry.py | 4 + vllm/models/deepseek_v4/__init__.py | 10 + .../deepseek_v4/common/mm_preprocess.py | 512 ++++++++++++++++++ vllm/models/deepseek_v4/common/vision.py | 232 ++++++++ vllm/models/deepseek_v4/nvidia/vl_model.py | 333 ++++++++++++ vllm/models/deepseek_v4/vl_stub.py | 14 + 6 files changed, 1105 insertions(+) create mode 100644 vllm/models/deepseek_v4/common/mm_preprocess.py create mode 100644 vllm/models/deepseek_v4/common/vision.py create mode 100644 vllm/models/deepseek_v4/nvidia/vl_model.py create mode 100644 vllm/models/deepseek_v4/vl_stub.py diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index db3464a4461d..5510e1b58636 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -363,6 +363,10 @@ "Cohere2VisionForConditionalGeneration", ), "Cosmos3ForConditionalGeneration": ("cosmos3", "Cosmos3ForConditionalGeneration"), + "DeepseekV4ForConditionalGeneration": ( + "vllm.models.deepseek_v4", + "DeepseekV4ForConditionalGeneration", + ), "Cosmos3EdgeForConditionalGeneration": ( "cosmos3_edge", "Cosmos3EdgeForConditionalGeneration", diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index c7dd948f79f8..7346b509e6ff 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -20,20 +20,30 @@ ) from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP + from .vl_stub import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) elif current_platform.is_xpu(): from .xpu.dspark import DSparkDeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] + from .vl_stub import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) else: from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, ) from .nvidia.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .nvidia.mtp import DeepSeekV4MTP # type: ignore[assignment] + from .nvidia.vl_model import ( # type: ignore[assignment] + DeepseekV4ForConditionalGeneration, + ) __all__ = [ "DSparkDeepseekV4ForCausalLM", "DeepSeekV4MTP", "DeepseekV4FP8Config", "DeepseekV4ForCausalLM", + "DeepseekV4ForConditionalGeneration", ] diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py new file mode 100644 index 000000000000..919146c551a4 --- /dev/null +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -0,0 +1,512 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Multimodal preprocessing for the DeepSeek-V4 vision variants +(DeepSeek-V4-Flash-Vision-Exp). + +The image transform and sentinel-block construction are ported from the +official repository's ``image_processor.py`` so that token counts bit-match +the reference. Each ``<｜deepseek_image｜>`` placeholder in the prompt expands +to a variable-length block of sentinel tokens; only positions with +``type == IMAGE`` receive vision embeddings, the other sentinels are looked +up from learned embedding vectors in the model. + +Unlike the reference (which uses out-of-vocab ids ``vocab_size + type``), +the sentinel block borrows five consecutive reserved tokenizer tokens +(``<|place_holder_mm_span_0431|>`` .. ``_0435|>``): they are special tokens +the tokenizer never emits from plain text, so the ids stay in-vocabulary and +work with stock token validation, logprobs and detokenization. The ids are +pure markers — every sentinel position's embedding is overwritten with +vision/sentinel vectors before the decoder layers see it, exactly like the +reference's out-of-vocab scheme. +""" + +import math +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import numpy as np +import torch +from PIL import Image, ImageOps +from transformers import BatchFeature +from typing_extensions import assert_never + +from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.parse import ImageSize, MultiModalDataItems +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + BaseProcessingInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.multimodal.processing.processor import ( + MultiModalPromptUpdates, + MultiModalPromptUpdatesApplyResult, + PlaceholderFeaturesInfo, + UpdateMode, + _plan_prompt_updates, +) +from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config + +IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5) +COMPRESS_PAD_TO = 4 + +IMAGE_PLACEHOLDER = "<｜deepseek_image｜>" + +# Sentinel roles borrow five consecutive ``<|place_holder_mm_span_XXXX|>`` +# tokens (reserved special tokens, never emitted from plain text). The order +# must match IMAGE_START..IMAGE_END above so that +# ``id == IMAGE_SENTINEL_BASE_ID + type``. +IMAGE_SENTINEL_BASE_ID = 129257 +IMAGE_SENTINEL_TOKEN_NAMES = ( + "<|place_holder_mm_span_0431|>", # IMAGE_START + "<|place_holder_mm_span_0432|>", # IMAGE_PAD + "<|place_holder_mm_span_0433|>", # IMAGE + "<|place_holder_mm_span_0434|>", # IMAGE_NEW_LINE + "<|place_holder_mm_span_0435|>", # IMAGE_END +) + + +def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor: + """Boolean mask for image-block sentinel positions (in-vocab ids).""" + return (token_ids >= IMAGE_SENTINEL_BASE_ID) & ( + token_ids < IMAGE_SENTINEL_BASE_ID + len(IMAGE_SENTINEL_TOKEN_NAMES) + ) + + +def validate_image_sentinel_ids(tokenizer) -> None: + """Check the borrowed sentinel ids against the tokenizer.""" + for i, name in enumerate(IMAGE_SENTINEL_TOKEN_NAMES): + token_id = tokenizer.convert_tokens_to_ids(name) + if token_id != IMAGE_SENTINEL_BASE_ID + i: + raise ValueError( + f"Image sentinel token {name!r} has id {token_id}, expected " + f"{IMAGE_SENTINEL_BASE_ID + i}; the DeepSeek-V4 vision " + "sentinel block requires these consecutive reserved ids." + ) + + +def grid_tokens(best_height, best_width, patch_size, downsample_ratio): + """Number of LLM tokens the aligner grid occupies (N-layout, including + row/align padding).""" + n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio) + n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio) + num_tokens = n_llm_h * (n_llm_w + 1) + 2 + if n_llm_h % 2 == 1: + num_tokens += n_llm_w + 1 + num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 + return n_llm_h, n_llm_w, num_tokens + + +def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): + r = height / width + max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 + max_h_float = max_w_float * r + if max_w_float < 1.0: + max_w = 1 + max_h = (max_n_token - 2) // (max_w + 1) + if max_h % 2 == 1: + max_h -= 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + elif max_h_float < 2.0: + max_h = 2 + max_w = ((max_n_token - 2) // max_h) - 1 + assert max_w > 1 + best_width = max_w * patch_size * downsample_ratio + best_height = max_h * patch_size * downsample_ratio + else: + max_w = math.floor(max_w_float) + max_h = math.floor(max_h_float) + if max_h % 2 == 1: + max_h -= 1 + beta = min( + max_w * patch_size * downsample_ratio / width, + max_h * patch_size * downsample_ratio / height, + ) + best_width = math.floor(width * beta / patch_size) * patch_size + best_height = math.floor(height * beta / patch_size) * patch_size + n_llm_h, n_llm_w, num_tokens = grid_tokens( + best_height, best_width, patch_size, downsample_ratio + ) + return n_llm_h, n_llm_w, best_height, best_width, num_tokens + + +def safe_resize( + height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token +): + max_n_token -= COMPRESS_PAD_TO - 1 + n_llm_h, n_llm_w, num_tokens = grid_tokens( + best_height, best_width, patch_size, downsample_ratio + ) + budget = max_n_token + while num_tokens > max_n_token: + n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio( + height, width, patch_size, downsample_ratio, budget + ) + budget -= 1 + return n_llm_h, n_llm_w, best_height, best_width + + +def load_image( + image: Image.Image, + *, + patch_size: int, + downsample_ratio: int, + max_n_token: int, + min_pixels: int, + max_wh_ratio: float | None, +): + """Transform one PIL image into ViT patches. + + Same math as the reference ``load_image``, except the image is already + decoded (vLLM supplies PIL images instead of a record dict). + """ + p = patch_size + image = image.convert("RGB") + width, height = image.size + if max_wh_ratio is not None and width > height * max_wh_ratio: + width = height * max_wh_ratio + if 0 < width * height < min_pixels: + ratio = (min_pixels / (width * height)) ** 0.5 + width = int(width * ratio) + height = int(height * ratio) + best_width = math.ceil(width / p) * p + best_height = math.ceil(height / p) * p + n_llm_h, n_llm_w, best_height, best_width = safe_resize( + height, width, best_height, best_width, p, downsample_ratio, max_n_token + ) + n_vit_h, n_vit_w = best_height // p, best_width // p + if max_wh_ratio is not None and image.width >= max_wh_ratio * image.height: + image = image.resize((best_width, best_height)) + else: + image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127)) + x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255 + x = ((x - 0.5) / 0.5).to(torch.bfloat16) + patches = ( + x.reshape(3, n_vit_h, p, n_vit_w, p) + .permute(1, 3, 0, 2, 4) + .reshape(n_vit_h * n_vit_w, 3, p, p) + ) + return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w + + +def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int): + """Builds the N-layout token types (final order) and the aligner-row order + for IMAGE slots.""" + compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO + pad_h = n_llm_h % 2 + rows = n_llm_h + pad_h + row_len = n_llm_w + 1 + pad_last = rows // 2 * row_len % 2 * 2 + types = torch.tensor( + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + + [IMAGE_PAD] * (row_len * pad_h), + dtype=torch.int64, + ) + order = torch.arange(rows * row_len).view(rows // 2, 2, row_len) + order = order.transpose(1, 2).reshape(-1) + image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64) + image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange( + n_llm_h * n_llm_w + ).view(n_llm_h, n_llm_w) + perm = image_idx[order] + perm = perm[perm >= 0] + types = torch.cat( + [ + torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_START]), + types[order], + torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64), + torch.tensor([IMAGE_END]), + ] + ) + return types, perm + + +def build_image_block_pad_free(n_llm_h: int, n_llm_w: int): + """``build_image_block`` without the position-dependent compressor pad. + + ``start_pos`` is chosen so that ``compress_pad == 0``; the pad is instead + prepended when the block is spliced into the final prompt, where its + position is known. + """ + return build_image_block(n_llm_h, n_llm_w, COMPRESS_PAD_TO - 1) + + +class DeepseekV4VLImageProcessor: + """Per-image transform (the PIL-input equivalent of the reference + ``load_image``).""" + + def __init__(self, config: DeepseekV4Config) -> None: + super().__init__() + self.patch_size = config.vision_patch_size + self.downsample_ratio = config.vision_downsample_ratio + self.max_n_token = config.vision_max_n_token + self.min_pixels = config.vision_min_pixels + self.max_wh_ratio = config.vision_max_wh_ratio + + def __call__(self, image: Image.Image): + return load_image( + image, + patch_size=self.patch_size, + downsample_ratio=self.downsample_ratio, + max_n_token=self.max_n_token, + min_pixels=self.min_pixels, + max_wh_ratio=self.max_wh_ratio, + ) + + +class DeepseekV4VLProcessor: + """Minimal stand-in for the HF processor of DeepSeek-V4 vision models. + + The official repository ships image preprocessing as plain functions in + ``image_processor.py`` (no ``auto_map`` processor), so this class wraps + their ports directly and the model loads without ``--trust-remote-code``. + + ``__call__`` returns a ``BatchFeature`` with one entry per image + (flattened across images): + + - ``patches``: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` bf16 ViT patches. + - ``vit_grid``: ``(num_images, 2)`` int64 ``[n_vit_h, n_vit_w]``. + - ``llm_grid``: ``(num_images, 2)`` int64 ``[n_llm_h, n_llm_w]``. + - ``perm``: concatenated per-image ``(n_llm_h * n_llm_w,)`` int64 index + selecting aligner outputs into the final N-layout order. + - ``types``: concatenated per-image pad-free sentinel block types; + ``block ids = IMAGE_SENTINEL_BASE_ID + types``. + """ + + def __init__(self, config: DeepseekV4Config) -> None: + super().__init__() + self.config = config + self.image_processor = DeepseekV4VLImageProcessor(config) + + def __call__( + self, + text: str | None = None, + images: Sequence[Image.Image] | None = None, + return_tensors: str | None = None, + **kwargs: Any, + ) -> BatchFeature: + patches_list = [] + vit_grid = [] + llm_grid = [] + perm_list = [] + types_list = [] + for image in images or []: + patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = self.image_processor(image) + types, perm = build_image_block_pad_free(n_llm_h, n_llm_w) + patches_list.append(patches) + vit_grid.append((n_vit_h, n_vit_w)) + llm_grid.append((n_llm_h, n_llm_w)) + perm_list.append(perm) + types_list.append(types) + + if not patches_list: + return BatchFeature({}) + + return BatchFeature( + { + "patches": torch.cat(patches_list), + "vit_grid": torch.tensor(vit_grid, dtype=torch.int64), + "llm_grid": torch.tensor(llm_grid, dtype=torch.int64), + "perm": torch.cat(perm_list), + "types": torch.cat(types_list), + } + ) + + +class DeepseekV4VLProcessingInfo(BaseProcessingInfo): + def get_hf_config(self) -> DeepseekV4Config: + return self.ctx.get_hf_config(DeepseekV4Config) + + def get_hf_processor(self, **kwargs: object) -> DeepseekV4VLProcessor: + if kwargs: + raise ValueError(f"Unexpected processor kwargs: {sorted(kwargs)}") + return DeepseekV4VLProcessor(self.get_hf_config()) + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + # ``safe_resize`` reserves COMPRESS_PAD_TO - 1 tokens of + # vision_max_n_token for the compressor-alignment pad, so the full + # sentinel block is bounded by vision_max_n_token; the margin is kept + # in case that reservation changes. + return { + "image": self.get_hf_config().vision_max_n_token + COMPRESS_PAD_TO - 1, + } + + def get_image_placeholder_token_id(self) -> int: + token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER) + if token_id is None: + raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}") + return token_id + + def get_image_size_with_most_features(self) -> ImageSize: + hf_config = self.get_hf_config() + patch_size = hf_config.vision_patch_size + downsample_ratio = hf_config.vision_downsample_ratio + # A square maximizes the ViT patch count (area) within the token + # budget; solve the budget-derived size directly to keep the dummy + # image small. + budget = hf_config.vision_max_n_token - (COMPRESS_PAD_TO - 1) + side = budget * patch_size * downsample_ratio + _, _, best_h, best_w, _ = solve_resize_ratio( + side, side, patch_size, downsample_ratio, budget + ) + return ImageSize(width=best_w, height=best_h) + + +class DeepseekV4VLDummyInputsBuilder( + BaseDummyInputsBuilder[DeepseekV4VLProcessingInfo] +): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + return IMAGE_PLACEHOLDER * mm_counts.get("image", 0) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions], + ) -> MultiModalDataDict: + size = self.info.get_image_size_with_most_features() + return { + "image": self._get_dummy_images( + width=size.width, + height=size.height, + num_images=mm_counts.get("image", 0), + overrides=cast(ImageDummyOptions | None, mm_options.get("image")), + ), + } + + +class DeepseekV4VLMultiModalProcessor( + BaseMultiModalProcessor[DeepseekV4VLProcessingInfo] +): + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + vit_grid = hf_inputs.get("vit_grid") + llm_grid = hf_inputs.get("llm_grid") + + if vit_grid is None or llm_grid is None: + empty = torch.empty(0, dtype=torch.long) + patch_sizes = perm_sizes = types_sizes = empty + else: + patch_sizes = vit_grid.prod(-1) + perm_sizes = llm_grid.prod(-1) + n_llm_h, n_llm_w = llm_grid[:, 0], llm_grid[:, 1] + # Pad-free block length; same formula as ``grid_tokens`` given + # the LLM grid. + types_sizes = ( + n_llm_h * (n_llm_w + 1) + + 2 + + (n_llm_h % 2) * (n_llm_w + 1) + + (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 + ) + + return { + "patches": MultiModalFieldConfig.flat_from_sizes("image", patch_sizes), + "vit_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), + "llm_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), + "perm": MultiModalFieldConfig.flat_from_sizes("image", perm_sizes), + "types": MultiModalFieldConfig.flat_from_sizes( + "image", types_sizes, keep_on_cpu=True + ), + } + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + image_token_id = self.info.get_image_placeholder_token_id() + validate_image_sentinel_ids(self.info.get_tokenizer()) + image_embed_id = IMAGE_SENTINEL_BASE_ID + IMAGE + + def get_image_replacement(item_idx: int) -> PromptUpdateDetails: + types: torch.Tensor = out_mm_kwargs["image"][item_idx]["types"].data + full = (IMAGE_SENTINEL_BASE_ID + types).tolist() + return PromptUpdateDetails.select_token_id(full, image_embed_id) + + return [ + PromptReplacement( + modality="image", + target=[image_token_id], + replacement=get_image_replacement, + ), + ] + + def _apply_token_matches_with_placeholders( + self, + token_ids: list[int], + mm_prompt_updates: MultiModalPromptUpdates, + ) -> tuple[ + list[int], + MultiModalPromptUpdatesApplyResult, + Mapping[str, list[PlaceholderFeaturesInfo]], + ]: + # Same as the base implementation, except that each image block gets + # its compressor-alignment pad (``3 - start_pos % 4`` IMAGE_PAD + # sentinels) prepended while splicing: the pad depends on the block's + # final position in the prompt, which is unknown when the (cacheable) + # prompt updates are built. + matched_updates, result = _plan_prompt_updates(token_ids, mm_prompt_updates) + placeholders: dict[str, list[PlaceholderFeaturesInfo]] = { + modality: [] for modality in mm_prompt_updates + } + + pad_id = IMAGE_SENTINEL_BASE_ID + IMAGE_PAD + + new_token_ids = list[int]() + prev_end_idx = 0 + for matched_update in matched_updates: + update = matched_update.update + match = matched_update.match + + if update.mode == UpdateMode.INSERT: + end_idx_to_insert = match.end_idx + elif update.mode == UpdateMode.REPLACE: + end_idx_to_insert = match.start_idx + else: + assert_never(update.mode) + + new_token_ids.extend(token_ids[prev_end_idx:end_idx_to_insert]) + start_idx = len(new_token_ids) + + tokens = list(update.content.full) + if tokens and update.modality == "image": + compress_pad = COMPRESS_PAD_TO - 1 - start_idx % COMPRESS_PAD_TO + tokens = [pad_id] * compress_pad + tokens + + if tokens: + content_is_embed = update.content.is_embed + is_embed = ( + content_is_embed(tokens) if content_is_embed is not None else None + ) + placeholders[update.modality].append( + PlaceholderFeaturesInfo( + modality=update.modality, + item_idx=update.item_idx, + start_idx=start_idx, + tokens=tokens, + is_embed=is_embed, + ) + ) + new_token_ids.extend(tokens) + + prev_end_idx = match.end_idx + + new_token_ids.extend(token_ids[prev_end_idx:]) + return new_token_ids, result, placeholders diff --git a/vllm/models/deepseek_v4/common/vision.py b/vllm/models/deepseek_v4/common/vision.py new file mode 100644 index 000000000000..a54c35e52c85 --- /dev/null +++ b/vllm/models/deepseek_v4/common/vision.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision tower (ViT + aligner) with TP-sharded linears. + +Ported from the official reference implementation +(deepseek-ai/DeepSeek-V4-Flash-Vision-Exp). Weight names match the HF +checkpoint so no renaming is needed at load time. Attention and MLP weights +are tensor-parallel sharded (replicated when the vision head count is not +divisible by TP size, or under ``--mm-encoder-tp-mode data``); the patch +embed and norms are replicated, so the residual stream is full-width on +every rank. +""" + +from functools import lru_cache + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention.mm_encoder_attention import ( + MMEncoderAttention, +) +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.models.vision import is_vit_use_data_parallel + + +@lru_cache(8) +def get_vision_cos_sin( + n_h: int, n_w: int, dim: int, theta: float +) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w) + wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w) + freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() + freqs = (freqs * inv_freq).flatten(1) + return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) + + +def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x1, x2 = x.float().chunk(2, dim=-1) + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype) + + +class DeepseekV4RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) + return (self.weight * x).to(dtype) + + +class DeepseekV4PatchEmbed(nn.Module): + def __init__(self, config): + super().__init__() + # Replicated: the residual stream is full-width on every rank. + self.proj = ReplicatedLinear( + 3 * config.vision_patch_size**2, + config.vision_dim, + bias=True, + quant_config=None, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out, _ = self.proj(x.flatten(1)) + return out + + +class DeepseekV4VisionAttention(nn.Module): + def __init__(self, config, prefix: str = ""): + super().__init__() + # Convention from minimax_m3/siglip: compute the TP/DP choice locally + # instead of threading a flag through every constructor. + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) + self.tp_size = ( + 1 if use_data_parallel else get_tensor_model_parallel_world_size() + ) + self.n_heads = config.vision_n_heads // self.tp_size + self.head_dim = config.vision_dim // config.vision_n_heads + self.hidden_size = config.vision_dim + self.wqkv = QKVParallelLinear( + config.vision_dim, + self.head_dim, + config.vision_n_heads, + bias=True, + quant_config=None, + prefix=f"{prefix}.wqkv", + disable_tp=use_data_parallel, + ) + self.wo = RowParallelLinear( + config.vision_dim, + config.vision_dim, + bias=True, + quant_config=None, + prefix=f"{prefix}.wo", + disable_tp=use_data_parallel, + ) + self.attn = MMEncoderAttention( + num_heads=self.n_heads, + head_size=self.head_dim, + prefix=f"{prefix}.attn", + ) + + def forward( + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + n = x.size(0) + qkv, _ = self.wqkv(x) + q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1)) + q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d) + k = apply_rotary(k, cos, sin).unsqueeze(0) + # One image per call: a dense batch, no varlen packing metadata. + o = self.attn(q, k, v.unsqueeze(0)) + out, _ = self.wo(o.reshape(n, -1)) + return out + + +class DeepseekV4VisionMLP(nn.Module): + def __init__(self, config, prefix: str = ""): + super().__init__() + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) + self.w1 = MergedColumnParallelLinear( + config.vision_dim, + [config.vision_inter_dim] * 2, + bias=False, + quant_config=None, + prefix=f"{prefix}.w1", + disable_tp=use_data_parallel, + ) + self.w2 = RowParallelLinear( + config.vision_inter_dim, + config.vision_dim, + bias=False, + quant_config=None, + prefix=f"{prefix}.w2", + disable_tp=use_data_parallel, + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.w1(x) + out, _ = self.w2(self.act_fn(gate_up)) + return out + + +class DeepseekV4VisionBlock(nn.Module): + def __init__(self, config, prefix: str = ""): + super().__init__() + self.norm1 = DeepseekV4RMSNorm(config.vision_dim) + self.attn = DeepseekV4VisionAttention(config, prefix=f"{prefix}.attn") + self.norm2 = DeepseekV4RMSNorm(config.vision_dim) + self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp") + + def forward( + self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + x = x + self.attn(self.norm1(x), cos, sin) + return x + self.mlp(self.norm2(x)) + + +class DeepseekV4ViT(nn.Module): + """DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE.""" + + def __init__(self, config): + super().__init__() + self.rope_dim = config.vision_dim // config.vision_n_heads // 2 + self.rope_theta = config.vision_rope_theta + self.patch_embed = DeepseekV4PatchEmbed(config) + self.blocks = nn.ModuleList( + [ + DeepseekV4VisionBlock(config, prefix=f"blocks.{i}") + for i in range(config.vision_n_layers) + ] + ) + self.norm = DeepseekV4RMSNorm(config.vision_dim) + + def forward( + self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int + ) -> torch.Tensor: + x = self.patch_embed(patches) + cos, sin = get_vision_cos_sin(n_vit_h, n_vit_w, self.rope_dim, self.rope_theta) + cos = cos.to(device=x.device) + sin = sin.to(device=x.device) + for block in self.blocks: + x = block(x, cos, sin) + return self.norm(x) + + +class DeepseekV4Aligner(nn.Module): + """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.""" + + def __init__(self, config): + super().__init__() + use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) + self.downsample_ratio = config.vision_downsample_ratio + in_dim = config.vision_dim * self.downsample_ratio**2 + self.w1 = ColumnParallelLinear( + in_dim, + config.hidden_size, + bias=True, + quant_config=None, + disable_tp=use_data_parallel, + ) + self.w2 = RowParallelLinear( + config.hidden_size, + config.hidden_size, + bias=True, + quant_config=None, + disable_tp=use_data_parallel, + ) + + def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: + r = self.downsample_ratio + x = x.view(n_vit_h, n_vit_w, -1).permute(2, 0, 1) + x = F.pad(x, (0, -n_vit_w % r, 0, -n_vit_h % r)) + x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1) + hidden, _ = self.w1(x) + out, _ = self.w2(F.gelu(hidden)) + return out diff --git a/vllm/models/deepseek_v4/nvidia/vl_model.py b/vllm/models/deepseek_v4/nvidia/vl_model.py new file mode 100644 index 000000000000..bd40b7bce9e7 --- /dev/null +++ b/vllm/models/deepseek_v4/nvidia/vl_model.py @@ -0,0 +1,333 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek-V4 vision variant (e.g. DeepSeek-V4-Flash-Vision-Exp). + +Thin multimodal wrapper around the text-only ``DeepseekV4ForCausalLM``: + +- ``vision`` ViT + ``aligner`` produce per-image embeddings for the IMAGE + sentinel positions; four learned vectors (``image_start`` / ``image_pad`` / + ``image_newline`` / ``image_end``) fill the remaining sentinel positions. +- Image placeholders (``<｜deepseek_image｜>``) are expanded by the processor + in ``common/mm_preprocess.py`` into sentinel blocks borrowing reserved + in-vocab tokens ``<|place_holder_mm_span_0431|>``..``_0435|>`` + (see ``common/vision.py`` for the tower itself). +- Merged embeddings enter the text model via ``inputs_embeds``, i.e. before + its hyper-connection stream expansion. Raw ``input_ids`` still flow into + the model so the MoE router can apply ``bias_vl`` to image tokens + (``requires_raw_input_tokens``). +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsEagle3, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + +from ..common.mm_preprocess import ( + IMAGE_PLACEHOLDER, + IMAGE_SENTINEL_BASE_ID, + DeepseekV4VLDummyInputsBuilder, + DeepseekV4VLMultiModalProcessor, + DeepseekV4VLProcessingInfo, + image_sentinel_mask, +) +from ..common.vision import DeepseekV4Aligner, DeepseekV4ViT +from .model import _make_deepseek_v4_weights_mapper + + +def _make_deepseek_v4_vl_weights_mapper( + expert_dtype: str, image_enabled: bool +) -> WeightsMapper: + """Text-checkpoint mapping rules re-rooted under ``language_model.``.""" + base = _make_deepseek_v4_weights_mapper(expert_dtype) + orig_to_new_prefix: dict[str, str | None] = { + "layers.": "language_model.model.layers.", + "embed.": "language_model.model.embed.", + "norm.": "language_model.model.norm.", + "hc_head": "language_model.model.hc_head", + "mtp.": "language_model.model.mtp.", + } + if not image_enabled: + orig_to_new_prefix.update({"vision.": None, "aligner.": None, "image_": None}) + return WeightsMapper( + orig_to_new_prefix=orig_to_new_prefix, + orig_to_new_regex=base.orig_to_new_regex, + orig_to_new_suffix={ + "head.weight": "language_model.lm_head.weight", + "embed.weight": "embed_tokens.weight", + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", + }, + orig_to_new_substr={ + ".shared_experts.w2": ".shared_experts.down_proj", + # The MTP/DSpark draft heads are not supported for the vision + # variant; drop their weights. + "mtp.": None, + }, + ) + + +@MULTIMODAL_REGISTRY.register_processor( + DeepseekV4VLMultiModalProcessor, + info=DeepseekV4VLProcessingInfo, + dummy_inputs=DeepseekV4VLDummyInputsBuilder, +) +class DeepseekV4ForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3 +): + """Multimodal entry point for DeepSeek-V4 checkpoints with a vision tower. + + ``SupportsEagle3`` (aux hidden-state plumbing for MTP/DSpark drafters) + delegates through ``language_model`` via the protocol defaults. + """ + + # The MoE router needs raw token ids to detect image sentinel tokens + # (borrowed reserved ids, see common/mm_preprocess.py) and apply bias_vl. + requires_raw_input_tokens = True + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality == "image": + return IMAGE_PLACEHOLDER + raise ValueError(f"Unsupported modality: {modality!r}") + + def __init__(self, *, vllm_config, prefix: str = "") -> None: + super().__init__() + model_config = vllm_config.model_config + config = model_config.hf_config + self.config = config + self.multimodal_config = model_config.multimodal_config + assert self.multimodal_config is not None + + image_enabled = ( + config.vision_n_layers > 0 + and self.multimodal_config.get_limit_per_prompt("image") > 0 + ) + with self._mark_tower_model(vllm_config, {"image"}): + self.vision: DeepseekV4ViT | None = None + self.aligner: DeepseekV4Aligner | None = None + self.image_start: nn.Parameter | None = None + self.image_end: nn.Parameter | None = None + self.image_newline: nn.Parameter | None = None + self.image_pad: nn.Parameter | None = None + if image_enabled: + self.vision = DeepseekV4ViT(config) + self.aligner = DeepseekV4Aligner(config) + for name in ( + "image_start", + "image_end", + "image_newline", + "image_pad", + ): + setattr( + self, + name, + nn.Parameter( + torch.empty(config.hidden_size, dtype=torch.float32) + ), + ) + self.vision.to(dtype=model_config.dtype) + self.aligner.to(dtype=model_config.dtype) + + with self._mark_language_model(vllm_config): + # The arch convertor routes any config with a vision tower to + # this wrapper class; mark the copy handed to the text backbone + # so it resolves to DeepseekV4ForCausalLM instead of recursing + # (with_hf_config deepcopies the config, the marker survives). + config._dsv4_vl_inner = True # type: ignore[attr-defined] + try: + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + hf_config=config, + prefix=maybe_prefix(prefix, "language_model"), + architectures=["DeepseekV4ForCausalLM"], + ) + finally: + del config._dsv4_vl_inner # type: ignore[attr-defined] + # The outer mapper (see load_weights) fully resolves HF names into + # this wrapper's namespace before AutoWeightsLoader strips the + # "language_model." prefix and delegates to the child's load_weights, + # so the child's own mapper must be a no-op. Its suffix rules are not + # idempotent (e.g. "lm_head.weight".endswith("head.weight") would + # re-fire "head.weight" -> "lm_head.weight"). + self.language_model.hf_to_vllm_mapper = WeightsMapper() + self.make_empty_intermediate_tensors = ( # type: ignore[method-assign] + self.language_model.make_empty_intermediate_tensors + ) + + expert_dtype = getattr(config, "expert_dtype", "fp4") + self.hf_to_vllm_mapper = _make_deepseek_v4_vl_weights_mapper( + expert_dtype, image_enabled + ) + + def _parse_and_validate_image_input(self, **kwargs: object) -> dict | None: + patches = kwargs.pop("patches", None) + if patches is None: + return None + vit_grid = kwargs.pop("vit_grid", None) + llm_grid = kwargs.pop("llm_grid", None) + perm = kwargs.pop("perm", None) + assert vit_grid is not None and llm_grid is not None and perm is not None + return { + "patches": patches, + "vit_grid": vit_grid, + "llm_grid": llm_grid, + "perm": perm, + } + + def _encode_image( + self, + patches: torch.Tensor, + n_vit_h: int, + n_vit_w: int, + perm: torch.Tensor, + ) -> torch.Tensor: + assert self.vision is not None and self.aligner is not None + image_embeds = self.aligner( + self.vision(patches, n_vit_h, n_vit_w), n_vit_h, n_vit_w + ) + # Reorder into the N-layout block order used in the prompt. + return image_embeds[perm.to(image_embeds.device)] + + def _process_image_input( + self, + patches: torch.Tensor, + vit_grid: torch.Tensor, + llm_grid: torch.Tensor, + perm: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + assert self.vision is not None and self.aligner is not None + patches = patches.to(self.aligner.w1.weight.dtype) + + embeds: list[torch.Tensor] = [] + vit_offset = 0 + llm_offset = 0 + for (n_vit_h, n_vit_w), (n_llm_h, n_llm_w) in zip( + vit_grid.tolist(), llm_grid.tolist(), strict=True + ): + n_vit = n_vit_h * n_vit_w + n_llm = n_llm_h * n_llm_w + embeds.append( + self._encode_image( + patches[vit_offset : vit_offset + n_vit], + n_vit_h, + n_vit_w, + perm[llm_offset : llm_offset + n_llm], + ) + ) + vit_offset += n_vit + llm_offset += n_llm + return tuple(embeds) + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None or self.vision is None: + return [] + return self._process_image_input( + image_input["patches"], + image_input["vit_grid"], + image_input["llm_grid"], + image_input["perm"], + ) + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + from vllm.model_executor.models.utils import _merge_multimodal_embeddings + + # All ids are in-vocab here: image-block sentinels are borrowed + # reserved tokens (their embedding rows are always overwritten below). + inputs_embeds = self.language_model.embed_input_ids(input_ids) + + if self.image_start is not None: + # Branch-free sentinel overwrite: safe inside compiled/captured + # regions (no data-dependent control flow). + sentinel_mask = image_sentinel_mask(input_ids) + if is_multimodal is not None: + # IMAGE positions get vision embeddings via the merge below. + sentinel_mask = sentinel_mask & ~is_multimodal.to(input_ids.device) + table = torch.stack( + [ + self.image_start, + self.image_pad, + self.image_pad, + self.image_newline, + self.image_end, + ] + ).to(inputs_embeds.dtype) + idx = (input_ids - IMAGE_SENTINEL_BASE_ID).clamp(0, 4) + inputs_embeds = torch.where( + sentinel_mask.unsqueeze(-1), table[idx], inputs_embeds + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + assert is_multimodal is not None + return _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors=None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.language_model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.language_model.compute_logits_local(hidden_states) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.language_model.get_expert_mapping() + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer for the MTP/DSpark draft model.""" + return self.language_model.get_mtp_target_hidden_states() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Map HF names into this wrapper's namespace up front and sort, so + # the "language_model." group reaches the child loader as one + # contiguous block (AutoWeightsLoader delegates per contiguous group, + # and the child's load_weights finalizes fused expert weights, which + # must not run on a partially loaded model). + mapped = sorted(self.hf_to_vllm_mapper.apply(weights), key=lambda x: x[0]) + loader = AutoWeightsLoader(self) + loaded_params = loader.load_weights(mapped) + # The child's load_weights already ran its post-load finalization. + self._weights_finalized = True + return loaded_params + + def process_weights_after_loading(self) -> None: + # Model-level post-load hook (called by the loader after any load + # format). Under DummyModelLoader the child's load_weights — and + # hence its finalize step — is bypassed, so run it here instead. + if getattr(self, "_weights_finalized", False): + return + self.language_model.process_weights_after_loading() diff --git a/vllm/models/deepseek_v4/vl_stub.py b/vllm/models/deepseek_v4/vl_stub.py new file mode 100644 index 000000000000..0f15a82f0190 --- /dev/null +++ b/vllm/models/deepseek_v4/vl_stub.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stub for platforms where the DeepSeek-V4 vision variant is unsupported.""" + +from torch import nn + + +class DeepseekV4ForConditionalGeneration(nn.Module): + def __init__(self, *, vllm_config, prefix: str = ""): + super().__init__() + raise NotImplementedError( + "DeepSeek-V4 vision (DeepseekV4ForConditionalGeneration) is only " + "supported on NVIDIA GPUs for now." + ) From 5d1ff7f119702ca0dbf76c2111011056affc809c Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:07:19 +0000 Subject: [PATCH 03/15] deepseek_v4: bias_vl image-routing hunks (Path 3 vision port, step 2/4) Hand-ports the vision-only hunks from the pinned Vision-Exp head's MoE routing chain onto this fork, keeping the existing ampere_sparse SM80 attention backend untouched: - csrc/libtorch_stable/moe/{moe_ops.h,torch_bindings.cpp, topk_softplus_sqrt_kernels.cu}: add bias_vl/image_sentinel_lo to the topk_softplus_sqrt native op (fast hash-table kernel and the generic templated kernel), so image sentinel tokens route by bias_vl instead of the text correction bias or hash table. - vllm/_custom_ops.py: thread the two new params through the op wrapper. - vllm/model_executor/layers/fused_moe/router/dsv4_topk.py: extend the Triton fast-path kernel with an optional per-row image-bias override. - vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py, router_factory.py, vllm/model_executor/layers/fused_moe/layer.py: plumb bias_vl/image_sentinel_lo through the router and FusedMoEFactory. - vllm/models/deepseek_v4/nvidia/model.py: create gate.bias_vl and image_sentinel_lo on vision checkpoints, pass them into FusedMoEFactory, require input_ids when bias_vl routing is active. - vllm/v1/worker/gpu/model_runner.py: keep input_ids on non-first PP ranks when the model sets requires_raw_input_tokens. Deliberately not ported: the pinned head's unrelated removal of the SM90-only dsv3_router_gemm kernel; MegaMoE shared-expert/sequence- parallel refactors (SM90+-only); XPU-only fallback paths. Source: PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX, branch codex/path3-vision-on-sm80-fork, commit 93df7ea6. CPU checks: py_compile and pyflakes clean on all touched Python files. CUDA/C++ files not compile-checked on this host (no build toolchain here); scoped to the same pattern as the existing topk_softplus_sqrt kernel and verified against the evidence repo's tracked source. --- csrc/libtorch_stable/moe/moe_ops.h | 4 +- .../moe/topk_softplus_sqrt_kernels.cu | 199 +++++++++++++++--- csrc/libtorch_stable/moe/torch_bindings.cpp | 3 +- vllm/_custom_ops.py | 4 + vllm/model_executor/layers/fused_moe/layer.py | 7 + .../layers/fused_moe/router/dsv4_topk.py | 31 +++ .../router/fused_topk_bias_router.py | 75 ++++++- .../layers/fused_moe/router/router_factory.py | 5 + vllm/models/deepseek_v4/nvidia/model.py | 28 ++- vllm/v1/worker/gpu/model_runner.py | 12 +- 10 files changed, 328 insertions(+), 40 deletions(-) diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 03525a05f9df..c0e35b90adaa 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -28,7 +28,9 @@ void topk_softplus_sqrt( const std::optional& correction_bias, const std::optional& input_ids, const std::optional& tid2eid, - const std::optional& is_padding); + const std::optional& is_padding, + const std::optional& bias_vl, + int64_t image_sentinel_lo); void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output, std::optional topk_ids, diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index 976eebac3ddd..6d6285bc978f 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -81,19 +81,92 @@ __launch_bounds__(128) __global__ int num_experts, float routed_scaling_factor, const HashIndType* input_ids, const HashIndType* tid2eid, - const bool* is_padding) { + const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int lane = threadIdx.x % 32; if (warp >= num_rows) return; const int64_t token_id = load_index_as_int64(input_ids, warp); const bool is_pad_row = is_padding != nullptr && is_padding[warp]; + // Image tokens carry five consecutive in-vocab sentinel ids starting at + // image_sentinel_lo (0 = disabled); they skip the tid2eid lookup and + // select experts by score + bias_vl instead. Ids above the sentinel block + // are regular special tokens and must not match. + const bool is_image = bias_vl != nullptr && image_sentinel_lo > 0 && + token_id >= image_sentinel_lo && + token_id < image_sentinel_lo + 5; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaGridDependencySynchronize(); #endif int expert = 0; float weight = 0.f; - if (lane < 6 && !is_pad_row) { + if (is_image && !is_pad_row) { + // Top-6 by (sqrtsoftplus(x) + bias_vl) over the full row; the routing + // weights are the raw sqrtsoftplus scores of the selected experts. + constexpr int MAX_PER_LANE = 12; // supports num_experts <= 384 + const int per_lane = num_experts / 32; + float vals[MAX_PER_LANE]; + float sels[MAX_PER_LANE]; + #pragma unroll + for (int i = 0; i < MAX_PER_LANE; ++i) { + if (i < per_lane) { + const int e = lane + 32 * i; + const float x = input[warp * num_experts + e]; + float v = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x)))); + if (isnan(v)) { + v = 0.f; + } + vals[i] = v; + sels[i] = v + bias_vl[e]; + } else { + vals[i] = 0.f; + sels[i] = -INFINITY; + } + } + int slot_expert[6]; + float slot_weight[6]; + #pragma unroll + for (int j = 0; j < 6; ++j) { + float best = -INFINITY; + int best_e = num_experts; + #pragma unroll + for (int i = 0; i < MAX_PER_LANE; ++i) { + if (sels[i] > best) { + best = sels[i]; + best_e = lane + 32 * i; + } + } + // Warp argmax reduce; ties go to the lower expert id. + #pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + const float other_best = VLLM_SHFL_XOR_SYNC(best, mask); + const int other_e = VLLM_SHFL_XOR_SYNC(best_e, mask); + if (other_best > best || (other_best == best && other_e < best_e)) { + best = other_best; + best_e = other_e; + } + } + float w = 0.f; + #pragma unroll + for (int i = 0; i < MAX_PER_LANE; ++i) { + if (lane + 32 * i == best_e) { + w = vals[i]; + sels[i] = -INFINITY; + } + } + #pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + w += VLLM_SHFL_XOR_SYNC(w, mask); + } + slot_expert[j] = best_e; + slot_weight[j] = w; + } + if (lane < 6) { + expert = slot_expert[lane]; + weight = slot_weight[lane]; + } + } else if (lane < 6 && !is_pad_row) { // only load and calculate for 6 experts expert = static_cast(tid2eid[token_id * 6 + lane]); const float x = input[warp * num_experts + expert]; @@ -127,7 +200,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, double routed_scaling_factor, const HashIndType* input_ids, const HashIndType* tid2eid, cudaStream_t stream, - const bool* is_padding) { + const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { if (num_rows == 0) return; auto* kernel = &dsv4HashTopkSoftplusSqrt; cudaLaunchConfig_t config = {}; @@ -141,7 +215,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, config.numAttrs = 1; const float scale = static_cast(routed_scaling_factor); cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows, - num_experts, scale, input_ids, tid2eid, is_padding); + num_experts, scale, input_ids, tid2eid, is_padding, + bias_vl, image_sentinel_lo); } #endif @@ -174,7 +249,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const HashIndType* input_ids, const HashIndType* tid2eid, - const bool* is_padding) { + const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { static_assert(std::is_same_v || std::is_same_v || std::is_same_v, @@ -323,9 +399,23 @@ __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); + // 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; + + // Hash MoE path: indices are predetermined from lookup table. Image rows + // (sentinel ids) fall through to the score-based top-k path with bias_vl. + const bool hash_row = USE_HASH && !use_vl_bias; + if (hash_row) { const int64_t token_expert_offset = token_id * static_cast(k); if (!is_pad_row) { #pragma unroll @@ -417,13 +507,13 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ if (isnan(val)) { val = 0.f; } - if (correction_bias) { + if (row_bias) { const int group_id = ii / ELTS_PER_LDG; const int local_id = ii % ELTS_PER_LDG; const int expert_idx = first_elt_read_by_thread + group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id; - val = val + correction_bias[expert_idx]; + val = val + row_bias[expert_idx]; } row_chunk[ii] = val; } @@ -488,8 +578,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // to global memory. (This will be a single) thread per row of the // input/output matrices. const int idx = k * thread_row + k_idx; - if (correction_bias != nullptr && should_process_row) { - max_val -= correction_bias[expert]; + if (row_bias != nullptr && should_process_row) { + max_val -= row_bias[expert]; } output[idx] = max_val; indices[idx] = @@ -576,7 +666,8 @@ void topkGatingSoftplusSqrtLauncherHelper( const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, - const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) { + const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding, + const float* bias_vl, int64_t image_sentinel_lo) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = @@ -605,12 +696,13 @@ void topkGatingSoftplusSqrtLauncherHelper( cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, is_padding); + input_ids, tid2eid, is_padding, bias_vl, + image_sentinel_lo); #else kernel<<>>( input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, is_padding); + input_ids, tid2eid, is_padding, bias_vl, image_sentinel_lo); #endif }) } @@ -624,7 +716,7 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \ - stream, is_padding); + stream, is_padding, bias_vl, image_sentinel_lo); #else #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -633,14 +725,14 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream, is_padding); \ + tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \ } else if (WARP_SIZE == 32) { \ topkGatingSoftplusSqrtLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream, is_padding); \ + tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -654,14 +746,16 @@ void topkGatingSoftplusSqrtKernelLauncher( const int topk, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, const HashIndType* tid2eid, - cudaStream_t stream, const bool* is_padding) { + cudaStream_t stream, const bool* is_padding, const float* bias_vl, + int64_t image_sentinel_lo) { #ifndef USE_ROCM if constexpr (std::is_same_v) { if (use_hash && topk == 6 && renormalize && (num_experts == 256 || num_experts == 384)) { launchDsv4HashTopk( gating_output, topk_weights, topk_indices, num_tokens, num_experts, - routed_scaling_factor, input_ids, tid2eid, stream, is_padding); + routed_scaling_factor, input_ids, tid2eid, stream, is_padding, + bias_vl, image_sentinel_lo); return; } } @@ -761,12 +855,30 @@ void dispatch_topk_softplus_sqrt_launch( const std::optional& correction_bias, const std::optional& input_ids, const std::optional& tid2eid, cudaStream_t stream, - const std::optional& is_padding) { + const std::optional& is_padding, + const std::optional& bias_vl, + int64_t image_sentinel_lo) { const float* bias_ptr = nullptr; if (correction_bias.has_value()) { bias_ptr = correction_bias.value().const_data_ptr(); } + const float* bias_vl_ptr = nullptr; + if (bias_vl.has_value()) { + const torch::stable::Tensor& bias_vl_tensor = bias_vl.value(); + STD_TORCH_CHECK(input_ids.has_value(), + "input_ids is required when bias_vl is set"); + STD_TORCH_CHECK( + bias_vl_tensor.scalar_type() == torch::headeronly::ScalarType::Float, + "bias_vl tensor must be float32"); + STD_TORCH_CHECK(bias_vl_tensor.dim() == 1, "bias_vl tensor must be 1D"); + STD_TORCH_CHECK(bias_vl_tensor.size(0) == num_experts, + "bias_vl size mismatch, expected: ", num_experts); + STD_TORCH_CHECK(bias_vl_tensor.is_contiguous(), + "bias_vl tensor must be contiguous"); + bias_vl_ptr = bias_vl_tensor.const_data_ptr(); + } + auto launch = [&](auto* topk_indices_ptr) { using OutIndType = typename std::remove_pointer::type; @@ -800,7 +912,8 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream, is_padding_ptr); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); } else { STD_TORCH_CHECK(tid2eid.value().scalar_type() == torch::headeronly::ScalarType::Int); @@ -810,7 +923,34 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream, is_padding_ptr); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); + } + } else if (bias_vl_ptr != nullptr) { + // Non-hash rows with bias_vl: image tokens (five sentinel ids starting + // at image_sentinel_lo) select experts by score + bias_vl. Dispatch the + // hash-index type on input_ids' dtype (no hash table is used here). + if (input_ids.value().scalar_type() == + torch::headeronly::ScalarType::Long) { + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, false, input_ids.value().const_data_ptr(), + static_cast(nullptr), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); + } else { + STD_TORCH_CHECK(input_ids.value().scalar_type() == + torch::headeronly::ScalarType::Int); + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + topk_indices_ptr, token_expert_indices.mutable_data_ptr(), + num_tokens, num_experts, topk, renormalize, routed_scaling_factor, + bias_ptr, false, input_ids.value().const_data_ptr(), + static_cast(nullptr), stream, is_padding_ptr, + bias_vl_ptr, image_sentinel_lo); } } else { vllm::moe::topkGatingSoftplusSqrtKernelLauncher( @@ -818,7 +958,8 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, false, static_cast(nullptr), - static_cast(nullptr), stream, is_padding_ptr); + static_cast(nullptr), stream, is_padding_ptr, + nullptr, 0); } }; @@ -843,7 +984,9 @@ void topk_softplus_sqrt( const std::optional& correction_bias, const std::optional& input_ids, const std::optional& tid2eid, - const std::optional& is_padding) { + const std::optional& is_padding, + const std::optional& bias_vl, + int64_t image_sentinel_lo) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); @@ -857,21 +1000,21 @@ void topk_softplus_sqrt( gating_output.const_data_ptr(), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, input_ids, tid2eid, stream, - is_padding); + is_padding, bias_vl, image_sentinel_lo); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_softplus_sqrt_launch<__half>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream, is_padding); + input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream, is_padding); + input_ids, tid2eid, stream, is_padding, bias_vl, image_sentinel_lo); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index c6467f1c1ff1..03d2b0ad18f4 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -20,7 +20,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, float " "routed_scaling_factor, Tensor? " - "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()"); + "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding, " + "Tensor? bias_vl=None, int image_sentinel_lo=0) -> ()"); // Calculate the result of moe by summing up the partial results // from all selected experts. topk_ids/expert_map are optional and, when diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 357e882834eb..7ebccf3b3c99 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2435,6 +2435,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, @@ -2447,6 +2449,8 @@ def topk_hash_softplus_sqrt( input_tokens, hash_indices_table, is_padding, + bias_vl, + image_sentinel_lo, ) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index cfbadba7aced..0dacb95219c8 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -140,6 +140,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, @@ -205,6 +207,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 disables vision routing (Deepseek V4) 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) @@ -319,6 +324,8 @@ def FusedMoEFactory( zero_expert_type=zero_expert_type, num_logical_experts=logical_num_experts, hash_indices_table=hash_indices_table, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) if params_dtype is None: diff --git a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py index dd861d4ebcf7..9f36a8d3edef 100644 --- a/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py +++ b/vllm/model_executor/layers/fused_moe/router/dsv4_topk.py @@ -44,9 +44,13 @@ def _dsv4_topk_kernel( 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, launch_pdl: tl.constexpr, + HAS_VL: tl.constexpr, ): row = tl.program_id(0) expert_offsets = tl.arange(0, BLOCK_N) @@ -58,6 +62,20 @@ def _dsv4_topk_kernel( if launch_pdl: tl.extra.cuda.gdc_wait() + 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) + logits = tl.load( gating_output_ptr + row * NUM_EXPERTS + expert_offsets, mask=expert_mask, @@ -101,11 +119,20 @@ 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 shape = (num_tokens, _TOPK) topk_weights = gating_output.new_empty(shape, dtype=torch.float32) topk_ids = gating_output.new_empty(shape, dtype=indices_dtype) + 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() if num_tokens > 0: _dsv4_topk_kernel[(num_tokens,)]( gating_output, @@ -113,9 +140,13 @@ 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), num_warps=1, launch_pdl=current_platform.is_arch_support_pdl(), + HAS_VL=has_vl, ) return topk_weights, topk_ids diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index bc019a58e2e2..fb8660d247ca 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -128,6 +128,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, ...]: from vllm.platforms import current_platform @@ -155,6 +157,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 @@ -184,6 +188,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 @@ -192,18 +198,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( @@ -211,6 +227,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() @@ -262,6 +281,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}") @@ -317,6 +338,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] @@ -332,10 +355,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)[ @@ -368,6 +420,8 @@ def __init__( 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, @@ -385,6 +439,11 @@ def __init__( # ``shared_expert_weight``, AFTER the routed top-k is renormalized. self.num_fused_shared_experts = num_fused_shared_experts self.shared_expert_weight = shared_expert_weight + # 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 @property def routing_method_type(self) -> RoutingMethodType: @@ -419,6 +478,8 @@ def _compute_routing( input_tokens=input_ids, hash_indices_table=self._hash_indices_table, routed_scaling_factor=self.routed_scaling_factor, + bias_vl=self.bias_vl.data if self.bias_vl is not None else None, + image_sentinel_lo=self.image_sentinel_lo, ) if self.num_fused_shared_experts > 0: diff --git a/vllm/model_executor/layers/fused_moe/router/router_factory.py b/vllm/model_executor/layers/fused_moe/router/router_factory.py index ff4874b20e61..d80fa840f8f6 100644 --- a/vllm/model_executor/layers/fused_moe/router/router_factory.py +++ b/vllm/model_executor/layers/fused_moe/router/router_factory.py @@ -60,6 +60,9 @@ def create_fused_moe_router( zero_expert_type: str | None = None, num_logical_experts: int | None = None, hash_indices_table: torch.Tensor | None = None, + # Deepseek V4 vision routing bias parameters + bias_vl: torch.Tensor | None = None, + image_sentinel_lo: int = 0, ) -> FusedMoERouter: """ Factory function to create the appropriate FusedMoERouter subclass based on @@ -211,6 +214,8 @@ def advertised_routing_method(groups: int | None) -> RoutingMethodType: hash_indices_table=hash_indices_table, num_fused_shared_experts=num_fused_shared_experts, shared_expert_weight=shared_expert_weight, + bias_vl=bias_vl, + image_sentinel_lo=image_sentinel_lo, ) if ( diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index db460dc8d270..ae8132f8615b 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -66,6 +66,7 @@ ) from vllm.model_executor.utils import set_weight_attrs from vllm.models.deepseek_v4.attention import DeepseekV4Attention +from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import ( DeepseekV4FlashInferMLAAttention, @@ -562,6 +563,12 @@ def __init__( 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: @@ -577,12 +584,25 @@ def __init__( ), 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: @@ -689,6 +709,8 @@ def _init_fused_moe_experts( router_logits_dtype=torch.float32, enable_eplb=parallel_config.enable_eplb, num_redundant_experts=eplb_config.num_redundant_experts, + bias_vl=getattr(self.gate, "bias_vl", None), + image_sentinel_lo=self.image_sentinel_lo, ) def forward( @@ -696,6 +718,10 @@ def forward( ) -> 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) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0f35c32afed7..d7fd8f6c5be7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1375,8 +1375,16 @@ def execute_model( **self.model_state.prepare_inputs(input_batch, self.req_states), } if not self.is_first_pp_rank: - # Update for non-first PP ranks. - model_inputs["input_ids"] = None + # Update for non-first PP ranks. Hidden states arrive via + # intermediate_tensors, so embeddings are never needed here. + # Raw input_ids are a separate concern: models that set + # requires_raw_input_tokens (e.g. DeepSeek V4 Vision's MoE + # image-token routing) thread input_ids through every decoder + # layer on every PP rank, not just to build the initial + # embedding on rank 0. Only null input_ids here when the model + # does not need it, so non-vision behaviour is unchanged. + if not self.model.requires_raw_input_tokens: + model_inputs["input_ids"] = None model_inputs["inputs_embeds"] = None # Prepare the intermediate tensors. From a21474b1a600ecd11230acd175cb0a3627368ec1 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:07:40 +0000 Subject: [PATCH 04/15] deepseek_v4: in-image bidirectional SWA window (Path 3 vision port, step 3/4) Widens combine_topk_swa_indices / build_flashinfer_mixed_sparse_indices and their Triton kernels with optional left_visible / right_visible / max_image_tokens parameters, so prefill tokens inside an image span get a bidirectionally widened sliding window instead of plain causal attention. Image-free batches stay byte-identical to prior behavior. Adds self.max_image_tokens on DeepseekV4Attention, computed from vision_max_n_token / vision_n_layers (zero on text-only checkpoints), which the widened-SWA hunk reads from the attention layer. Source: PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX, branch codex/path3-vision-on-sm80-fork, commits b6721bf1 (cache_utils.py) and 24d07d22 (attention.py). CPU checks: py_compile and pyflakes clean on both touched files. --- vllm/models/deepseek_v4/attention.py | 7 + .../deepseek_v4/common/ops/cache_utils.py | 149 ++++++++++++++---- 2 files changed, 123 insertions(+), 33 deletions(-) diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 4f246ecaa0e0..eadf1bc1a1a1 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -204,6 +204,13 @@ def __init__( 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 diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index 3fc33775b661..cb228c72b813 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -535,10 +535,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 ) @@ -567,6 +573,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 @@ -627,6 +636,7 @@ class CompileKey: TOP_K: int COMPRESS_RATIO: int WINDOW_SIZE: int + IMAGE_WIDTH: int PADDED_TOP_K: int input_variant: TritonPointerInputVariant @@ -648,11 +658,14 @@ def kernel( 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) @@ -681,7 +694,19 @@ def kernel( 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 @@ -694,18 +719,19 @@ def kernel( 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) @@ -721,6 +747,7 @@ def dispatch( # type: ignore[override] 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( @@ -733,6 +760,7 @@ def dispatch( # type: ignore[override] TOP_K=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=WINDOW_SIZE, + IMAGE_WIDTH=image_width, PADDED_TOP_K=padded_topk, input_variant=input_variant, ) @@ -742,10 +770,19 @@ def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: 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: @@ -762,11 +799,14 @@ def compile(self, compile_key: CompileKey) -> None: 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), ) @@ -785,8 +825,13 @@ def __call__( 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), @@ -796,11 +841,14 @@ def __call__( 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]), ) @@ -827,19 +875,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, window_size + - padded_topk]`` (the first ``window_size`` 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. + 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 - assert decode_swa_indices.shape[-1] == window_size + 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 @@ -888,7 +950,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, window_size + padded_topk), + (num_tokens, swa_total_width + padded_topk), dtype=torch.int32, device=decode_swa_indices.device, ) @@ -898,7 +960,7 @@ def build_flashinfer_mixed_sparse_indices( if num_tokens == 0: return sparse_indices, sparse_topk_lens - window_block_size = triton.next_power_of_2(max(window_size, 1)) + window_block_size = triton.next_power_of_2(max(swa_index_width, 1)) topk_block_size = triton.next_power_of_2(max(padded_topk, 1)) max_block_size = max(window_block_size, topk_block_size) num_warps = 4 if max_block_size >= 256 else 1 @@ -922,6 +984,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, @@ -935,6 +999,8 @@ def build_flashinfer_mixed_sparse_indices( compressed_span, 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, @@ -989,6 +1055,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, @@ -1002,6 +1070,8 @@ def _build_flashinfer_mixed_sparse_indices_kernel( compressed_block_span, 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, @@ -1013,14 +1083,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, WINDOW_SIZE, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < WINDOW_SIZE + 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) @@ -1063,7 +1137,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + WINDOW_SIZE + + SWA_TOTAL_WIDTH + offset, values, mask=mask, @@ -1077,7 +1151,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, WINDOW_SIZE + compressed_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + compressed_len) return prefill_idx = token_idx - NUM_DECODE_TOKENS @@ -1089,13 +1163,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, WINDOW_SIZE, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_TOTAL_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < WINDOW_SIZE + mask = offset < SWA_TOTAL_WIDTH pos_offset = swa_start_pos + offset block_indices = pos_offset // swa_block_size block_numbers = tl.load( @@ -1139,10 +1222,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + WINDOW_SIZE + + SWA_TOTAL_WIDTH + offset, slot_ids, mask=mask, ) - tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + topk_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_TOTAL_WIDTH + topk_len) From 7de2c14c4b7d54533c95150ab5c85dc0156e9552 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:08:13 +0000 Subject: [PATCH 05/15] deepseek_v4: three of the five vision boot fixes (Path 3 vision port, step 4/4) Three of the five boot errors found bringing the ported vision model up on 4x CMP 170HX (the other two are a launch-flag change, not a source fix, and the load_weights finalize fix already committed separately): - vllm/multimodal/processing/processor.py: add _plan_prompt_updates, a thin variant of _apply_matches that returns the ordered (update, match) pairs instead of applying them, so mm_preprocess.py can splice in compressor-alignment padding whose size depends on each match's final position in the prompt. - vllm/models/deepseek_v4/common/mm_preprocess.py: override _call_hf_processor on DeepseekV4VLMultiModalProcessor to tokenize the prompt and merge input_ids into the processor output (fixes `KeyError: 'input_ids'`); override _hf_processor_applies_updates to return False so the normal placeholder-expansion path runs instead of being skipped (fixes "found 0 prompt placeholders"). - vllm/v1/worker/gpu/cudagraph_utils.py: add the same requires_raw_input_tokens gate to create_forward_fn's capture path that gpu/model_runner.py's regular forward-input-prep path already had, so CUDA-graph warmup on non-first PP ranks stops crashing with "DeepSeek V4 vision MoE routing requires input_ids." Default behavior is unchanged for every other model. Source: PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX, branch codex/path3-vision-on-sm80-fork, commits 7ab7cf2f (processor.py), 0c991df3 (mm_preprocess.py), c1e2de67 (cudagraph_utils.py). See that repo's docs/VISION-PORT.md for the full five-fix account, including fix 1 (drop --safetensors-load-strategy eager; a launch-flag change, no source diff) and fix 4 (load_weights finalize call, already in the "vision-only files + registry wiring" ancestor commit's follow-up). CPU checks: py_compile and pyflakes clean on all three touched files. --- .../deepseek_v4/common/mm_preprocess.py | 49 +++++++++++++++ vllm/multimodal/processing/processor.py | 63 +++++++++++++++++++ vllm/v1/worker/gpu/cudagraph_utils.py | 17 ++++- 3 files changed, 127 insertions(+), 2 deletions(-) diff --git a/vllm/models/deepseek_v4/common/mm_preprocess.py b/vllm/models/deepseek_v4/common/mm_preprocess.py index 919146c551a4..997518b7c73e 100644 --- a/vllm/models/deepseek_v4/common/mm_preprocess.py +++ b/vllm/models/deepseek_v4/common/mm_preprocess.py @@ -391,6 +391,55 @@ def get_dummy_mm_data( class DeepseekV4VLMultiModalProcessor( BaseMultiModalProcessor[DeepseekV4VLProcessingInfo] ): + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + # ``DeepseekV4VLProcessor`` (our stand-in "HF processor") only knows + # how to turn PIL images into ViT patches; unlike a real HF + # processor it never tokenizes ``text`` and so never populates + # ``input_ids``. The base ``_apply_hf_processor_text_mm`` requires + # ``input_ids`` in the returned `BatchFeature` (it pops the key + # right after calling us), so tokenize the prompt ourselves and + # merge it in. The prompt still carries the literal, un-expanded + # ``IMAGE_PLACEHOLDER`` text/token once per image at this point; + # expansion into the full per-image sentinel block happens later, + # in ``_get_prompt_updates`` / ``_apply_token_matches_with_placeholders``. + processed_data = super()._call_hf_processor( + prompt=prompt, + mm_data=mm_data, + mm_kwargs=mm_kwargs, + tok_kwargs=tok_kwargs, + ) + + if "input_ids" not in processed_data: + tokenizer = self.info.get_tokenizer() + encode_kwargs = dict(tok_kwargs) + encode_kwargs.setdefault("add_special_tokens", False) + input_ids = tokenizer.encode(prompt, **encode_kwargs) + processed_data["input_ids"] = [input_ids] + + return processed_data + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + # ``DeepseekV4VLProcessor`` never expands the raw image placeholder + # token into the per-image sentinel block itself (it only produces + # pixel-derived fields); that expansion is done afterwards by + # ``_get_prompt_updates`` / ``_apply_token_matches_with_placeholders``. + # The base class default (``True`` whenever mm data is present) + # would tell vLLM to skip that step, leaving the raw placeholder + # token in the prompt and failing placeholder validation. + return False + def _get_mm_fields_config( self, hf_inputs: BatchFeature, diff --git a/vllm/multimodal/processing/processor.py b/vllm/multimodal/processing/processor.py index f43b45ff7e4d..6800ad1ccfdb 100644 --- a/vllm/multimodal/processing/processor.py +++ b/vllm/multimodal/processing/processor.py @@ -828,6 +828,69 @@ def _apply_matches( return cast(list[_S], out_seqs), out_result +class _PlannedPromptUpdate(NamedTuple): + """One resolved update paired with where it matched in the prompt.""" + + update: "ResolvedPromptUpdate" + match: "PromptTargetMatch" + + +def _plan_prompt_updates( + prompt: _S, + mm_prompt_updates: "MultiModalPromptUpdates", + tokenizer: TokenizerLike | None = None, +) -> tuple[list[_PlannedPromptUpdate], "MultiModalPromptUpdatesApplyResult"]: + """ + Compatibility shim: resolve the matches that :func:`_apply_matches` would + apply, without applying them, so a caller can splice in extra content + (e.g. alignment padding) that depends on the final position of each + match. Returns a flat, left-to-right ordered list of + `_PlannedPromptUpdate` alongside the same result mapping that + `_apply_matches` returns. + """ + mm_item_counts = {m: len(items) for m, items in mm_prompt_updates.items()} + + out_result: MultiModalPromptUpdatesApplyResult = { + m: [None] * len(items) for m, items in mm_prompt_updates.items() + } + planned = list[_PlannedPromptUpdate]() + + mm_found_counts = { + m: sum(r is not None for r in res) for m, res in out_result.items() + } + if _all_items_found(mm_item_counts, mm_found_counts): + return planned, out_result + + prev_end_idx = 0 + while True: + mode, matches_to_apply = _find_matches( + prompt, + mm_prompt_updates, + tokenizer, + prev_end_idx=prev_end_idx, + current_result=out_result, + ) + + if mode is None: + break # No more matches to find + + for (modality, item_idx), (match, update_idx) in matches_to_apply: + matched_update = mm_prompt_updates[modality][item_idx][update_idx] + planned.append(_PlannedPromptUpdate(update=matched_update, match=match)) + out_result[modality][item_idx] = update_idx + + # Exclude overlapping matches + prev_end_idx = match.end_idx + + mm_found_counts = { + m: sum(r is not None for r in res) for m, res in out_result.items() + } + if _all_items_found(mm_item_counts, mm_found_counts): + break + + return planned, out_result + + def apply_token_matches( prompt: list[int], mm_prompt_updates: "MultiModalPromptUpdates", diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 3a174cba80d7..35c3eed10517 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -484,8 +484,21 @@ def create_forward_fn( **model_state.prepare_dummy_inputs(num_reqs, num_tokens), } if not self.is_first_pp_rank: - # Update for non-first PP ranks. - model_inputs["input_ids"] = None + # Update for non-first PP ranks. Hidden states arrive via + # intermediate_tensors, so embeddings are never needed here. + # Raw input_ids are a separate concern: models that set + # requires_raw_input_tokens (e.g. DeepSeek V4 Vision's MoE + # image-token routing) thread input_ids through every decoder + # layer on every PP rank, not just to build the initial + # embedding on rank 0. Only null input_ids here when the model + # does not need it, mirroring the same exception already + # applied in gpu/model_runner.py's non-capture forward path + # (model_runner.py's prepare-inputs branch for non-first PP + # ranks) -- this file's capture path was missing that gate, + # so CUDA-graph warmup crashed those models with "requires + # input_ids" instead of matching the real forward path. + if not getattr(model, "requires_raw_input_tokens", False): + model_inputs["input_ids"] = None model_inputs["inputs_embeds"] = None assert intermediate_tensors is not None model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] From ddae19b73043d10ce26e8cdbe8b0c96f115cd05a Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:08:41 +0000 Subject: [PATCH 06/15] deepseek_v4: fix hc_attn_fn_broadcast never finalized for real weight loads Boot fix 4 of 5 from the vision port: DeepseekV4ForCausalLM.load_weights did not call self.process_weights_after_loading() before returning. The outer VL wrapper (vl_model.py) sets self._weights_finalized = True right after delegating to this method, assuming finalization already ran here; without this call, the wrapper's own process_weights_after_loading() (the only other place finalize_mega_moe_weights / finalize_mhc_broadcast_weights could run) saw the flag already set and skipped them, so hc_attn_fn_broadcast was never populated for any real (non-dummy) weight load. This crashed DeepseekV4DecoderLayer.forward with `AssertionError: self.hc_attn_fn_broadcast is not None` during profile_run's image dummy-forward. Fix: call self.process_weights_after_loading() before returning, mirroring the pinned Vision head's (2c8af219) ordering exactly. Source: PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX, branch codex/path3-vision-on-sm80-fork, commit f7c0f5e7. CPU checks: py_compile and pyflakes clean. --- vllm/models/deepseek_v4/nvidia/model.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ae8132f8615b..88f16fb76787 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1541,7 +1541,23 @@ def get_mtp_target_hidden_states(self) -> torch.Tensor | None: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self, skip_substrs=["mtp.", "vision.", "aligner.", "image_", "_vl"]) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + loaded_params = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + # Mirror the pinned Vision head (2c8af219) ordering: self-invoke the + # post-load finalize step here so it always runs immediately after a + # real (non-dummy) weight load, before any forward pass can observe + # unfinalized state (e.g. hc_attn_fn_broadcast still None). + # + # This also matters for the outer VL wrapper + # (DeepseekV4ForConditionalGeneration.load_weights, vl_model.py): + # it sets `self._weights_finalized = True` right after delegating to + # this method, on the assumption that finalization already happened + # here. Without this call, that flag was set prematurely, and the + # wrapper's own process_weights_after_loading() (the only other + # place finalize_mega_moe_weights/finalize_mhc_broadcast_weights + # could run) then early-returned and skipped it entirely — so + # hc_attn_fn_broadcast was never populated for a real weight load. + self.process_weights_after_loading() + return loaded_params def process_weights_after_loading(self) -> None: # Model-level post-load hook: runs for every loader, including From c6e1956041a068e29138a53a4fabbfc6479f8410 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:09:33 +0000 Subject: [PATCH 07/15] patches: add safetensors F8_E8M0 overlay for DeepSeek-V4 ue8m0 scale tensors Adds patches/safetensors_torch_f8_e8m0.py, a full replacement for the installed safetensors/torch.py that adds an F8_E8M0 dtype entry (torch.float8_e8m0fnu, Torch 2.5.0+) to the load/save dtype tables. Without it, safetensors raises an unrecognized-dtype error on a checkpoint shard carrying a ue8m0 block-scale tensor. This is a third-party-package overlay, not a vLLM source patch, applied at image build time by copying it over the installed safetensors package (see patches/README.md and docker/Dockerfile.sm80). Source: PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX, patches/safetensors_torch.py (matches the copy baked into the published vision image, receipted at dsv4-vision-exp-4card-20260902T0440Z/receipts/safetensors_torch.py). CPU checks: py_compile and pyflakes clean. --- patches/README.md | 22 + patches/safetensors_torch_f8_e8m0.py | 597 +++++++++++++++++++++++++++ 2 files changed, 619 insertions(+) create mode 100644 patches/README.md create mode 100644 patches/safetensors_torch_f8_e8m0.py diff --git a/patches/README.md b/patches/README.md new file mode 100644 index 000000000000..e7abb63d5572 --- /dev/null +++ b/patches/README.md @@ -0,0 +1,22 @@ +# Runtime overlays + +Files here are not part of the vLLM Python package tree. The build copies +them over an installed third-party package inside the image, after +`pip install`, instead of patching vLLM source. + +## `safetensors_torch_f8_e8m0.py` + +Overlay for `safetensors/torch.py` (the installed `safetensors` package, +not vLLM). Adds an `F8_E8M0` dtype entry (`torch.float8_e8m0fnu`, Torch +2.5.0+) to the dtype tables that `safetensors` uses to load and save +tensors, so a checkpoint shard that stores a `ue8m0` block-scale tensor +loads without `safetensors` raising an unrecognized-dtype error. + +Applied in `docker/Dockerfile.sm80` with: + +```dockerfile +COPY patches/safetensors_torch_f8_e8m0.py /opt/venv/lib/python3.12/site-packages/safetensors/torch.py +``` + +Source: `PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX`, +`patches/safetensors_torch.py`. diff --git a/patches/safetensors_torch_f8_e8m0.py b/patches/safetensors_torch_f8_e8m0.py new file mode 100644 index 000000000000..8513881023e4 --- /dev/null +++ b/patches/safetensors_torch_f8_e8m0.py @@ -0,0 +1,597 @@ +import os +import sys +from collections import defaultdict +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +import torch +from safetensors import ( + TensorSpec, + deserialize, + safe_open, + serialize, + serialize_file, +) + + +def storage_ptr(tensor: torch.Tensor) -> int: + try: + return tensor.untyped_storage().data_ptr() + except Exception: + # Fallback for torch==1.10 + try: + return tensor.storage().data_ptr() + except NotImplementedError: + # Fallback for meta storage + return 0 + + +def _end_ptr(tensor: torch.Tensor) -> int: + if tensor.nelement(): + stop = tensor.view(-1)[-1].data_ptr() + _SIZE[tensor.dtype] + else: + stop = tensor.data_ptr() + return stop + + +def storage_size(tensor: torch.Tensor) -> int: + try: + return tensor.untyped_storage().nbytes() + except AttributeError: + # Fallback for torch==1.10 + try: + return tensor.storage().size() * _SIZE[tensor.dtype] + except NotImplementedError: + # Fallback for meta storage + # On torch >=2.0 this is the tensor size + return tensor.nelement() * _SIZE[tensor.dtype] + + +def _filter_shared_not_shared( + tensors: List[Set[str]], state_dict: Dict[str, torch.Tensor] +) -> List[Set[str]]: + filtered_tensors = [] + for shared in tensors: + if len(shared) < 2: + filtered_tensors.append(shared) + continue + + areas = [] + for name in shared: + tensor = state_dict[name] + areas.append((tensor.data_ptr(), _end_ptr(tensor), name)) + areas.sort() + + _, last_stop, last_name = areas[0] + filtered_tensors.append({last_name}) + for start, stop, name in areas[1:]: + if start >= last_stop: + filtered_tensors.append({name}) + else: + filtered_tensors[-1].add(name) + last_stop = stop + + return filtered_tensors + + +def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]: + tensors = defaultdict(set) + for k, v in state_dict.items(): + if ( + v.device != torch.device("meta") + and storage_ptr(v) != 0 + and storage_size(v) != 0 + ): + # Need to add device as key because of multiple GPU. + tensors[(v.device, storage_ptr(v), storage_size(v))].add(k) + tensors = list(sorted(tensors.values())) + tensors = _filter_shared_not_shared(tensors, state_dict) + return tensors + + +def _is_complete(tensor: torch.Tensor) -> bool: + return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _SIZE[ + tensor.dtype + ] == storage_size(tensor) + + +def _remove_duplicate_names( + state_dict: Dict[str, torch.Tensor], + *, + preferred_names: Optional[List[str]] = None, + discard_names: Optional[List[str]] = None, +) -> Dict[str, List[str]]: + if preferred_names is None: + preferred_names = [] + preferred_names = set(preferred_names) + if discard_names is None: + discard_names = [] + discard_names = set(discard_names) + + shareds = _find_shared_tensors(state_dict) + to_remove = defaultdict(list) + for shared in shareds: + complete_names = set( + [name for name in shared if _is_complete(state_dict[name])] + ) + if not complete_names: + raise RuntimeError( + "Error while trying to find names to remove to save state dict, but found no suitable name to keep" + f" for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model" + " since you could be storing much more memory than needed. Please refer to" + " https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an" + " issue." + ) + + keep_name = sorted(list(complete_names))[0] + + # Mechanism to preferentially select keys to keep + # coming from the on-disk file to allow + # loading models saved with a different choice + # of keep_name + preferred = complete_names.difference(discard_names) + if preferred: + keep_name = sorted(list(preferred))[0] + + if preferred_names: + preferred = preferred_names.intersection(complete_names) + if preferred: + keep_name = sorted(list(preferred))[0] + for name in sorted(shared): + if name != keep_name: + to_remove[keep_name].append(name) + return to_remove + + +def save_model( + model: torch.nn.Module, + filename: str, + metadata: Optional[Dict[str, str]] = None, + force_contiguous: bool = True, +): + """ + Saves a given torch model to specified filename. + This method exists specifically to avoid tensor sharing issues which are + not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors) + + Args: + model (`torch.nn.Module`): + The model to save on disk. + filename (`str`): + The filename location to save the file + metadata (`Dict[str, str]`, *optional*): + Extra information to save along with the file. + Some metadata will be added for each dropped tensors. + This information will not be enough to recover the entire + shared structure but might help understanding things + force_contiguous (`boolean`, *optional*, defaults to True): + Forcing the state_dict to be saved as contiguous tensors. + This has no effect on the correctness of the model, but it + could potentially change performance if the layout of the tensor + was chosen specifically for that reason. + """ + state_dict = model.state_dict() + to_removes = _remove_duplicate_names(state_dict) + + for kept_name, to_remove_group in to_removes.items(): + for to_remove in to_remove_group: + if metadata is None: + metadata = {} + + if to_remove not in metadata: + # Do not override user data + metadata[to_remove] = kept_name + del state_dict[to_remove] + if force_contiguous: + state_dict = {k: v.contiguous() for k, v in state_dict.items()} + try: + save_file(state_dict, filename, metadata=metadata) + except ValueError as e: + msg = str(e) + msg += " Or use save_model(..., force_contiguous=True), read the docs for potential caveats." + raise ValueError(msg) + + +def load_model( + model: torch.nn.Module, + filename: Union[str, os.PathLike], + strict: bool = True, + device: Union[str, int] = "cpu", + *, + backend: str = "mmap", +) -> Tuple[List[str], List[str]]: + """ + Loads a given filename onto a torch model. + This method exists specifically to avoid tensor sharing issues which are + not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors) + + Args: + model (`torch.nn.Module`): + The model to load onto. + filename (`str`, or `os.PathLike`): + The filename location to load the file from. + strict (`bool`, *optional*, defaults to True): + Whether to fail if you're missing keys or having unexpected ones. + When false, the function simply returns missing and unexpected names. + device (`Union[str, int]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular torch device locations. + backend (`str`, *optional*, defaults to `"mmap"`): + Storage backend used to serve tensor bytes. `"mmap"` (default) + and `"pread"` uses `pread(2)` to read tensor bytes. + + Returns: + `(missing, unexpected): (List[str], List[str])` + `missing` are names in the model which were not modified during loading + `unexpected` are names that are on the file, but weren't used during + the load. + """ + state_dict = load_file(filename, device=device, backend=backend) + model_state_dict = model.state_dict() + to_removes = _remove_duplicate_names( + model_state_dict, preferred_names=state_dict.keys() + ) + missing, unexpected = model.load_state_dict(state_dict, strict=False) + missing = set(missing) + for to_remove_group in to_removes.values(): + for to_remove in to_remove_group: + if to_remove not in missing: + unexpected.append(to_remove) + else: + missing.remove(to_remove) + if strict and (missing or unexpected): + missing_keys = ", ".join([f'"{k}"' for k in sorted(missing)]) + unexpected_keys = ", ".join([f'"{k}"' for k in sorted(unexpected)]) + error = f"Error(s) in loading state_dict for {model.__class__.__name__}:" + if missing: + error += f"\n Missing key(s) in state_dict: {missing_keys}" + if unexpected: + error += f"\n Unexpected key(s) in state_dict: {unexpected_keys}" + raise RuntimeError(error) + return missing, unexpected + + +def save( + tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None +) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, torch.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.torch import save + import torch + + tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + keep_references_alive = [] # to avoid garbage collection of temporary numpy arrays while we write to disk + serialized = serialize( + _flatten_as_ptr(tensors, keep_references_alive), metadata=metadata + ) + result = bytes(serialized) + return result + + +def save_file( + tensors: Dict[str, torch.Tensor], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +): + """ + Saves a dictionary of tensors into `filename` in safetensors format. + There is no mechanism in place to prevent the caller from modifying the data while a file save occurs, + please be wary when calling `save_file` and modifying tensors referenced in the `tensors` dict concurrently; + it may lead to corrupted files. + + Args: + tensors (`Dict[str, torch.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.torch import save_file + import torch + + tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + keep_references_alive = [] # to avoid garbage collection of temporary numpy arrays while we write to disk + serialize_file( + _flatten_as_ptr(tensors, keep_references_alive), filename, metadata=metadata + ) + + +def load_file( + filename: Union[str, os.PathLike], + device: Union[str, int] = "cpu", + *, + backend: str = "mmap", +) -> Dict[str, torch.Tensor]: + """ + Loads a safetensors file into torch format. + + Args: + filename (`str`, or `os.PathLike`): + The name of the file which contains the tensors + device (`Union[str, int]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular torch device locations. + backend (`str`, *optional*, defaults to `"mmap"`): + Storage backend used to serve tensor bytes. `"mmap"` (default) + and `"pread"` uses `pread(2)` to read tensor bytes. + + Returns: + `Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` + + Example: + + ```python + from safetensors.torch import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + with safe_open(filename, framework="pt", device=device, backend=backend) as f: + return f.get_tensors() + + +def load(data: bytes) -> Dict[str, torch.Tensor]: + """ + Loads a safetensors file into torch format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` on cpu + + Example: + + ```python + from safetensors.torch import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = deserialize(data) + return _view2torch(flat) + + +# torch.float8 formats require 2.1; we do not support these dtypes on earlier versions +_float8_e4m3fn = getattr(torch, "float8_e4m3fn", None) +_float8_e4m3fnuz = getattr(torch, "float8_e4m3fnuz", None) +_float8_e5m2 = getattr(torch, "float8_e5m2", None) +_float8_e5m2fnuz = getattr(torch, "float8_e5m2fnuz", None) +_float8_e8m0 = getattr(torch, "float8_e8m0fnu", None) +_float4_e2m1_x2 = getattr(torch, "float4_e2m1fn_x2", None) + +_SIZE = { + torch.int64: 8, + torch.float32: 4, + torch.int32: 4, + torch.bfloat16: 2, + torch.float16: 2, + torch.int16: 2, + torch.uint8: 1, + torch.int8: 1, + torch.bool: 1, + torch.float64: 8, + torch.complex64: 8, + _float8_e4m3fn: 1, + _float8_e4m3fnuz: 1, + _float8_e5m2: 1, + _float8_e5m2fnuz: 1, + _float8_e8m0: 1, + _float4_e2m1_x2: 1, +} + +if hasattr(torch, "uint64"): # Torch 2.3.0+ + _SIZE.update( + { + torch.uint64: 8, + torch.uint32: 4, + torch.uint16: 2, + } + ) + +_TYPES = { + "F64": torch.float64, + "F32": torch.float32, + "F16": torch.float16, + "BF16": torch.bfloat16, + "I64": torch.int64, + "I32": torch.int32, + "I16": torch.int16, + "I8": torch.int8, + "U8": torch.uint8, + "BOOL": torch.bool, + "F8_E4M3": _float8_e4m3fn, + "F8_E4M3FNUZ": _float8_e4m3fnuz, + "F8_E5M2": _float8_e5m2, + "F8_E5M2FNUZ": _float8_e5m2fnuz, + "C64": torch.complex64, +} + +if hasattr(torch, "uint64"): # Torch 2.3.0+ + _TYPES.update( + { + "U64": torch.uint64, + "U32": torch.uint32, + "U16": torch.uint16, + } + ) + +if hasattr(torch, "float8_e8m0fnu"): # Torch 2.5.0+: ue8m0 scale tensors + _TYPES.update( + { + "F8_E8M0": torch.float8_e8m0fnu, + } + ) + + +def _getdtype(dtype_str: str) -> torch.dtype: + return _TYPES[dtype_str] + + +def _view2torch(safeview) -> Dict[str, torch.Tensor]: + result = {} + for k, v in safeview: + dtype = _getdtype(v["dtype"]) + if len(v["data"]) == 0: + # Workaround because frombuffer doesn't accept zero-size tensors + assert any(x == 0 for x in v["shape"]) + arr = torch.empty(v["shape"], dtype=dtype) + else: + arr = torch.frombuffer(v["data"], dtype=dtype).reshape(v["shape"]) + if sys.byteorder == "big": + arr = torch.from_numpy(arr.numpy().byteswap(inplace=False)) + result[k] = arr + + return result + + +def _to_ndarray(tensor: torch.Tensor): + if tensor.device.type != "cpu": + # Moving tensor to cpu before saving + tensor = tensor.to("cpu") + + import ctypes + + import numpy as np + + # When shape is empty (scalar), np.prod returns a float + # we need a int for the following calculations + length = int(np.prod(tensor.shape).item()) + bytes_per_item = _SIZE[tensor.dtype] + + total_bytes = length * bytes_per_item + + ptr = tensor.data_ptr() + if ptr == 0: + return np.empty( + 0 + ), 0 # XXX: bogus value we don't really care if we return a tensor here + newptr = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_ubyte)) + data = np.ctypeslib.as_array(newptr, (total_bytes,)) # no internal copy + if sys.byteorder == "big": + NPDTYPES = { + torch.int64: np.int64, + torch.float32: np.float32, + torch.int32: np.int32, + # XXX: This is ok because both have the same width + torch.bfloat16: np.float16, + torch.float16: np.float16, + torch.int16: np.int16, + torch.uint8: np.uint8, + torch.int8: np.int8, + torch.bool: bool, + torch.float64: np.float64, + # XXX: This is ok because both have the same width and byteswap is a no-op anyway + _float8_e4m3fn: np.uint8, + _float8_e4m3fnuz: np.uint8, + _float8_e5m2: np.uint8, + _float8_e5m2fnuz: np.uint8, + _float8_e8m0: np.uint8, + _float4_e2m1_x2: np.uint8, + torch.complex64: np.complex64, + } + npdtype = NPDTYPES[tensor.dtype] + # Not in place as that would potentially modify a live running model + data = data.view(npdtype).byteswap(inplace=False) + return data, tensor + + +def _evaluate_tensors_for_save(tensors: Dict[str, torch.Tensor]) -> None: + if not isinstance(tensors, dict): + raise ValueError( + f"Expected a dict of [str, torch.Tensor] but received {type(tensors)}" + ) + + sparse_tensors = [] + for k, v in tensors.items(): + if not isinstance(v, torch.Tensor): + raise ValueError( + f"Key `{k}` is invalid, expected torch.Tensor but received {type(v)}" + ) + + if v.layout != torch.strided: + sparse_tensors.append(k) + + if sparse_tensors: + raise ValueError( + f"You are trying to save a sparse tensors: `{sparse_tensors}` which this library does not support." + " You can make it a dense tensor before saving with `.to_dense()` but be aware this might" + " make a much larger file than needed." + ) + + shared_pointers = _find_shared_tensors(tensors) + failing = [] + for names in shared_pointers: + if len(names) > 1: + failing.append(names) + + if failing: + raise RuntimeError( + f""" + Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: {failing}. + A potential way to correctly save your model is to use `save_model`. + More information at https://huggingface.co/docs/safetensors/torch_shared_tensors + """ + ) + + +def _flatten_as_ptr( + tensors: Dict[str, torch.Tensor], keep_alive_buffer: List +) -> Dict[str, Dict[str, Any]]: + _evaluate_tensors_for_save(tensors) + flattened = {} + for k, v in tensors.items(): + # XXX: doing this check later on instead of in _evaluate_tensors_for_save + # since on old versions of torch, SparseTensorImpl do not implement is_contiguous + # and we do the sparsity check in _evaluate_tensors_for_save. + if not v.is_contiguous(): + raise ValueError( + f"You are trying to save a non contiguous tensor: `{k}` which is not allowed. It either means you" + " are trying to save tensors which are reference of each other in which case it's recommended to save" + " only the full tensors, and reslice at load time, or simply call `.contiguous()` on your tensor to" + " pack it before saving." + ) + arr, tensor_ref = _to_ndarray(v) + keep_alive_buffer.append((arr, tensor_ref)) + flattened[k] = TensorSpec( + dtype=str(v.dtype).split(".")[-1], + shape=v.shape, + data_ptr=arr.ctypes.data, + data_len=arr.nbytes, + ) + return flattened From 5f1af1c2cdf35b5530b4ed32f1e0645024ed38b2 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:10:57 +0000 Subject: [PATCH 08/15] model_runner: zero stale CUDA-graph padding rows for vision input_ids (RCA Bug A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root-causes the c>=4 EngineCore/shm_broadcast crash on the vision build: Worker_PP3 dies with a device-side assert (NCCL watchdog), which EngineCore then reports downstream as a shm_broadcast "cancelled" timeout. Vision's bias_vl MoE routing (added in an earlier commit) made every PP rank read real input_ids on every decoder layer, but an unfilled CUDA-graph padding row can hold a stale, possibly sentinel-range token id left over from a prior, larger batch at the same bucket size. Text-only models never hit this because they keep input_ids=None on non-first ranks. Fix: zero the padding rows of input_ids right after the real rows are written, gated on requires_raw_input_tokens (false for every model except this vision port) — an SM80-fork-only code path, so this does not change behavior on SM90+. Source: seanphan/pixelml issue #79, RCA comment, "Minimal patch proposal (Bug A)". CPU checks: py_compile and pyflakes clean. --- vllm/v1/worker/gpu/model_runner.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d7fd8f6c5be7..ec8ec72f4139 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -948,6 +948,17 @@ def prepare_inputs( num_tokens = scheduler_output.total_num_scheduled_tokens num_tokens_after_padding = batch_desc.num_tokens assert num_tokens > 0 + if ( + num_tokens_after_padding > num_tokens + and getattr(self.model, "requires_raw_input_tokens", False) + ): + # Models that thread raw input_ids through every PP rank (vision + # bias_vl MoE routing) must never read a stale, possibly + # sentinel-range token id left in an unfilled CUDA-graph padding + # row from a prior, larger batch at this bucket size. Text-only + # models keep input_ids=None on non-first ranks and are + # unaffected (see cudagraph_utils.py's matching guard). + self.input_buffers.input_ids[num_tokens:num_tokens_after_padding].fill_(0) if envs.VLLM_MOE_SKIP_PADDING: # Mark trailing cudagraph-padding rows so kernels can skip work for # them when supported. From 3e5aabcaf41027287969c7d091d9794c79049641 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:10:57 +0000 Subject: [PATCH 09/15] dflash: clamp valid_ctx_end against out-of-bounds read on full rejection (RCA Bug B) Root-causes the 131k-token engine crash on the vision build: a heavily-chunked long prefill (--max-num-batched-tokens 2048) can produce a boundary chunk where DSpark verify rejects every scheduled token, collapsing valid_ctx_end (ctx_end - num_rejected) to ctx_start or below. The kernel then read before this request's own context window, raising "Triton Error [CUDA]: an illegal memory access was encountered" in _prepare_dflash_inputs_kernel. Fix: clamp valid_ctx_end to be at least ctx_start + 1, mirroring the existing upper-bound-clamp style already used in this kernel for block indices. Only changes behavior for the num_rejected >= num_ctx edge case; architecture-agnostic Triton index math, so it applies identically on SM90+, and only prevents an out-of-bounds read that would otherwise fault on any architecture. Source: seanphan/pixelml issue #79, RCA comment, "Minimal patch proposal (Bug B)". CPU checks: py_compile and pyflakes clean. --- vllm/v1/worker/gpu/spec_decode/dflash/speculator.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 333648637bd7..d132c950d7b4 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -521,6 +521,14 @@ def _prepare_dflash_inputs_kernel( num_rejected = tl.load(num_rejected_ptr + req_idx) valid_ctx_end = ctx_end - num_rejected + # A boundary-sized chunk (common on long, heavily-chunked prefills, + # e.g. 131k tokens at --max-num-batched-tokens 2048) can have every + # scheduled token rejected by DSpark verify, collapsing valid_ctx_end + # to ctx_start or below. Clamp so the position lookup below never reads + # before this request's own context window; downstream consumers only + # ever read the bonus-token/query rows this feeds, so a clamped, inert + # position for a fully-rejected chunk changes nothing observable. + valid_ctx_end = tl.maximum(valid_ctx_end, ctx_start + 1) num_sampled = tl.load(num_sampled_ptr + req_idx) if num_sampled > 0: From 7dee6a9ec32a3bfa99ce11e18d3cb9dd753e18bb Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 05:40:59 +0000 Subject: [PATCH 10/15] docs: SM80 fork provenance and README section Base f8ea5bb16 (haosdent/vllm dsv4-flash-a100 lineage, credit allover326/ Lasimeri for the underlying Ampere sparse-MLA technique). Flags the missing 8-file DSpark-under-PP patch diff as a follow-up item. (cherry picked from commit 5dc67e9c5caf9a62842a10426eabbed2e8b64943) --- README.md | 36 ++++++++++++++++++++++++++ docs/SM80.md | 72 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+) create mode 100644 docs/SM80.md diff --git a/README.md b/README.md index ca31a8cce94d..862855a9e6ef 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,42 @@ For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us. --- +## PixelML SM80 fork + +This fork carries NVIDIA Ampere (compute capability 8.0) fallback kernels +for model families whose upstream vLLM support targets Hopper or newer. +It exists for the PixelML `club-170hx` hardware group (4x NVIDIA CMP 170HX, +GA100 silicon, 64 GiB HBM2e per card, no FP4/FP8 tensor cores, no NVLink). +Full provenance, licenses, and per-branch status are in +[`docs/SM80.md`](docs/SM80.md). + +### Supported models + +| Model | Branch | Status | Notes | +| --- | --- | --- | --- | +| DeepSeek-V4-Flash (text) | `sm80` | Serving, published | `ghcr.io/pixelml/club-170hx:vllm-deepseek-v4-sm80-20260902`; see `club-170hx` `docs/MODEL-STATUS.md` | +| DeepSeek-V4-Flash-Vision-Exp | `sm80` (vision hunks in `PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX`) | Serving, published | Reverse-ported vision path onto this fork's attention backend; see that repo's `docs/VISION-PORT.md` | +| GLM-5.3-Flash | `glm53-sm80` | In progress (repo-only Phase A as of 2026-09-03) | See issue [#103](https://github.com/seanphan/pixelml/issues/103) for the port plan and inventory | + +### Build instructions + +```bash +git checkout sm80 # or glm53-sm80 once that branch has a bootable model +TORCH_CUDA_ARCH_LIST=8.0 docker build -f Dockerfile.fullbuild16 -t . +``` + +A from-source build takes about 60 minutes on the `club-170hx` build host. +See `club-170hx` `docs/DOCKER-IMAGE.md` for the exact published image +digests and launch commands. + +### Links + +- Platform manual and benchmarks: [`PixelML/club-170hx`](https://github.com/PixelML/club-170hx) +- DeepSeek-V4 vision port evidence: [`PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX`](https://github.com/PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX) +- GLM-5.3-Flash tracking issue: [seanphan/pixelml#103](https://github.com/seanphan/pixelml/issues/103) + +--- + ## About vLLM is a fast and easy-to-use library for LLM inference and serving. diff --git a/docs/SM80.md b/docs/SM80.md new file mode 100644 index 000000000000..2411e2e22521 --- /dev/null +++ b/docs/SM80.md @@ -0,0 +1,72 @@ +# SM80 (Ampere / CMP 170HX) fork: provenance and status + +This fork exists to run PixelML CMP 170HX (NVIDIA GA100 silicon, compute +capability 8.0, no FP4/FP8 tensor cores, no NVLink) on models whose upstream +vLLM support targets Hopper or newer. It is not a general-purpose vLLM fork; +it exists to carry Ampere fallback kernels that upstream does not (yet) ship. + +## Branch `sm80` — base + +- Base commit: `f8ea5bb163c161ef38b401d055cc5fd4a934091a`, + "[Attention] DeepSeek-V4 sparse MLA on SM8x (A100/A800)". +- Lineage: this commit descends from the `haosdent/vllm` fork + (`dsv4-flash-a100` line), which itself builds on prior Ampere DeepSeek-V4 + work credited to `allover326` and `Lasimeri` (Triton sparse-MLA kernel + patterns, manual e4m3 bit tricks). See + `PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX` `docs/VISION-PORT.md` and + the `club-170hx` evidence trail for the full history. +- Carries a working `vllm/models/deepseek_v4/ampere/ampere_sparse.py` SM80 + sparse-MLA attention backend, which the mainline `deepseek_v4` model does + not have (mainline's sparse-MLA backends declare support only for compute + capability 9/10/12). + +### Known gap in this commit + +The `club-170hx` published image (`ghcr.io/pixelml/club-170hx:vllm-deepseek-v4-sm80-20260902`) +was built from this base commit plus an additional 8-file patch set +(`.dockerignore`, `vllm/config/speculative.py`, +`vllm/model_executor/layers/sparse_attn_indexer.py`, +`vllm/models/deepseek_v4/nvidia/dspark.py`, +`vllm/models/deepseek_v4/nvidia/model.py`, +`vllm/v1/worker/gpu/model_runner.py`, `vllm/v1/worker/gpu/pp_utils.py`, +`vllm/v1/worker/gpu/spec_decode/dspark/utils.py`) that unlocks DSpark +speculative decoding under pipeline parallelism. The raw diff for that patch +set (`vllm-source-patches.diff`) was baked into the container image but was +not found, in this session, in any git commit or archived diff file +reachable from this host — the evidence repo's receipts folder kept the +build log and file list but not the diff blob itself. **Action item:** pull +that diff from the built image (`docker diff` against a clean +`f8ea5bb16` checkout, or from wherever the original build host still has +it) and commit it here as a second commit on `sm80` before this branch is +treated as the source of truth for future builds. Until then, `sm80` HEAD is +the base commit only, not yet the exact source of the published image. + +### Licensing + +- Upstream vLLM: Apache-2.0. No constraint. +- `haosdent/vllm` fork: no LICENSE file changes found; treated as + Apache-2.0-inherited (fork of an Apache-2.0 project, no relicensing + notice seen). +- `allover326/deepseek-v4-cmp170hx` packaging patterns: license not + confirmed from public pages as of 2026-09-03. Reused pattern + (PP + speculative-decoding unlock is a scheduler-level technique, not a + copied file) rather than verbatim code, pending confirmation. +- `Lasimeri/vllm-dsv4-ampere`: no LICENSE file. Treated as all-rights-reserved; + only the documented *technique* (Triton sparse-MLA kernel shape, e4m3 bit + manipulation) was reused, not copied source, per the prior-art review on + issue #103. + +## Branch `glm53-sm80` — GLM-5.3-Flash port + +Branched from `sm80`. See the issue #103 comment for the full port plan, +Hopper-dependency inventory, and hours estimate. Summary of the base +decision: neither applying vLLM PR #53906 (`glm5_next` support) directly +onto `sm80`, nor reverse-porting `sm80`'s fallbacks onto the PR head, +was the starting point. A third-party fork, `wtdcode/vllm-backport` +(Apache-2.0), already carries a working Ampere/SM80 build of `glm5next` +including a generic Triton sparse-MLA backend +(`vllm/v1/attention/backends/mla/triton_mla_sparse.py`) built on the +identical `ampere_sparse.py` file this fork already has. This branch vendors +those already-Ampere-adapted files from `wtdcode/vllm-backport` rather than +re-deriving Ampere fallbacks from scratch. Every vendored file is listed +with its wtdcode commit reference in the commit that adds it. From be954434a873fb139e24a778f6d8a8e5d0d71a18 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 05:52:05 +0000 Subject: [PATCH 11/15] docs: correct glm53-sm80 lineage claim (independent of sm80, not branched from it) (cherry picked from commit 588a947e77a717111eaa291d1b95e2bb1e0a7a51) --- docs/SM80.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/SM80.md b/docs/SM80.md index 2411e2e22521..d64aa1f771ae 100644 --- a/docs/SM80.md +++ b/docs/SM80.md @@ -58,7 +58,12 @@ the base commit only, not yet the exact source of the published image. ## Branch `glm53-sm80` — GLM-5.3-Flash port -Branched from `sm80`. See the issue #103 comment for the full port plan, +Base commit: `3bec275739` (wtdcode/vllm-backport `master` tip, +"scheduler: recover from invalid KV blocks on hybrid memory allocator"). +**Not** a descendant of `sm80` — wtdcode's GLM-5.3-Flash Ampere enablement +has an independent lineage from our DeepSeek-V4 sparse-MLA base. The two +branches share only the general SM80 sparse-MLA technique, not a common +git ancestor. See the issue #103 comment for the full port plan, Hopper-dependency inventory, and hours estimate. Summary of the base decision: neither applying vLLM PR #53906 (`glm5_next` support) directly onto `sm80`, nor reverse-porting `sm80`'s fallbacks onto the PR head, From 13b93329cfc01ff8d1d2c360cc23f7eb4b2be090 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:12:10 +0000 Subject: [PATCH 12/15] docs: reflect consolidation, add Qwen3.8-27B and known limits Updates docs/SM80.md and the README supported-models table now that sm80 carries the DSpark patch, the vision port, and the safetensors overlay as commits instead of a documented gap: - removes the "known gap" note (the DSpark 8-file patch is now committed on this branch, source recovered from the club-170hx image's build receipts); - adds the Qwen3.8-27B row (stock vLLM, no fork changes needed); - adds a known-limits section for the vision build's two crashes (c>=4, 131k-token prefill), both root-caused and patched on this branch per seanphan/pixelml#79, not yet re-verified on a GPU boot; - points the build instructions at docker/Dockerfile.sm80. --- README.md | 14 +++++++++++-- docs/SM80.md | 59 +++++++++++++++++++++++++++++++++++----------------- 2 files changed, 52 insertions(+), 21 deletions(-) diff --git a/README.md b/README.md index 862855a9e6ef..c185895369bc 100644 --- a/README.md +++ b/README.md @@ -33,14 +33,24 @@ Full provenance, licenses, and per-branch status are in | Model | Branch | Status | Notes | | --- | --- | --- | --- | | DeepSeek-V4-Flash (text) | `sm80` | Serving, published | `ghcr.io/pixelml/club-170hx:vllm-deepseek-v4-sm80-20260902`; see `club-170hx` `docs/MODEL-STATUS.md` | -| DeepSeek-V4-Flash-Vision-Exp | `sm80` (vision hunks in `PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX`) | Serving, published | Reverse-ported vision path onto this fork's attention backend; see that repo's `docs/VISION-PORT.md` | +| DeepSeek-V4-Flash-Vision-Exp | `sm80` | Serving, published; two known limits below | Reverse-ported vision path (Path 3), consolidated onto this branch; see `docs/SM80.md` and `PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX` `docs/VISION-PORT.md` | | GLM-5.3-Flash | `glm53-sm80` | In progress (repo-only Phase A as of 2026-09-03) | See issue [#103](https://github.com/seanphan/pixelml/issues/103) for the port plan and inventory | +| Qwen3.8-27B | stock vLLM (no fork changes needed) | Serving, published | No SM80 fallback required; see `club-170hx` `docs/MODEL-STATUS.md` | + +### Known limits (DeepSeek-V4-Flash-Vision-Exp) + +- **`c>=4` crash**: an EngineCore crash at concurrency 4 and above. + Root-caused and fixed on `sm80`; not yet re-verified on a GPU boot. + See [seanphan/pixelml#79](https://github.com/seanphan/pixelml/issues/79). +- **131k-token crash**: prefill past 65,000 tokens crashes the engine. + Root-caused and fixed on `sm80`; not yet re-verified on a GPU boot. + See [seanphan/pixelml#79](https://github.com/seanphan/pixelml/issues/79). ### Build instructions ```bash git checkout sm80 # or glm53-sm80 once that branch has a bootable model -TORCH_CUDA_ARCH_LIST=8.0 docker build -f Dockerfile.fullbuild16 -t . +TORCH_CUDA_ARCH_LIST=8.0 docker build -f docker/Dockerfile.sm80 -t . ``` A from-source build takes about 60 minutes on the `club-170hx` build host. diff --git a/docs/SM80.md b/docs/SM80.md index d64aa1f771ae..99f61e3668ec 100644 --- a/docs/SM80.md +++ b/docs/SM80.md @@ -20,26 +20,47 @@ it exists to carry Ampere fallback kernels that upstream does not (yet) ship. not have (mainline's sparse-MLA backends declare support only for compute capability 9/10/12). -### Known gap in this commit +### Consolidated onto this branch -The `club-170hx` published image (`ghcr.io/pixelml/club-170hx:vllm-deepseek-v4-sm80-20260902`) -was built from this base commit plus an additional 8-file patch set -(`.dockerignore`, `vllm/config/speculative.py`, -`vllm/model_executor/layers/sparse_attn_indexer.py`, -`vllm/models/deepseek_v4/nvidia/dspark.py`, -`vllm/models/deepseek_v4/nvidia/model.py`, -`vllm/v1/worker/gpu/model_runner.py`, `vllm/v1/worker/gpu/pp_utils.py`, -`vllm/v1/worker/gpu/spec_decode/dspark/utils.py`) that unlocks DSpark -speculative decoding under pipeline parallelism. The raw diff for that patch -set (`vllm-source-patches.diff`) was baked into the container image but was -not found, in this session, in any git commit or archived diff file -reachable from this host — the evidence repo's receipts folder kept the -build log and file list but not the diff blob itself. **Action item:** pull -that diff from the built image (`docker diff` against a clean -`f8ea5bb16` checkout, or from wherever the original build host still has -it) and commit it here as a second commit on `sm80` before this branch is -treated as the source of truth for future builds. Until then, `sm80` HEAD is -the base commit only, not yet the exact source of the published image. +`sm80` now carries, as separate commits on top of the base commit above: + +- the 8-file DSpark-under-pipeline-parallelism patch set used by the + published text image `ghcr.io/pixelml/club-170hx:vllm-deepseek-v4-sm80-20260902` + (recovered from that image's build receipts, + `dsv4-vision-exp-4card-20260902T0440Z/receipts/vllm-source-patches.diff`); +- the DeepSeek-V4-Flash-Vision-Exp reverse port (Path 3), four commits plus + a fifth for the `load_weights` finalize fix — see + `PixelML/DeepSeek-V4-Flash-Vision-Exp-CMP-170HX` `docs/VISION-PORT.md` + for the full account; +- the `safetensors` F8_E8M0 overlay patch needed to load a checkpoint + shard carrying a `ue8m0` block-scale tensor (`patches/`); +- two root-cause fixes from issue #79 for the vision build's two known + crashes (see "Known limits" below). + +`sm80` HEAD is now the exact source used to build +`docker/Dockerfile.sm80`; images no longer need bind-mounted fixes on +top of a pinned base. + +### Known limits (vision build) + +- **`c>=4` crash**: an EngineCore/`shm_broadcast` timeout at concurrency + 4 and above, root-caused to a stale CUDA-graph padding row and fixed + on this branch. See + [seanphan/pixelml#79](https://github.com/seanphan/pixelml/issues/79) + for the RCA and the patch commit on this branch + (`vllm/v1/worker/gpu/model_runner.py`, "RCA Bug A"). +- **131k-token crash**: prefill past 65,000 tokens crashed the engine + with a Triton illegal-memory-access in the DSpark drafter, root-caused + to an out-of-bounds read on a fully-rejected chunk and fixed on this + branch. See + [seanphan/pixelml#79](https://github.com/seanphan/pixelml/issues/79) + for the RCA and the patch commit on this branch + (`vllm/v1/worker/gpu/spec_decode/dflash/speculator.py`, "RCA Bug B"). +- Neither fix has a GPU boot verified against it yet on this branch as + of this commit; both patches are the evidence-backed primary + hypothesis from the issue #79 RCA, applied as static source changes. + Re-run the `c>=4` ladder and the 131k-token prefill test on the next + vision-server boot window to confirm. ### Licensing From 2f9ea97b337e91c38366f21adec17f4e62d43194 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 06:12:47 +0000 Subject: [PATCH 13/15] docker: add Dockerfile.sm80 for from-source SM80 builds Copies the fork's Dockerfile.fullbuild16 lineage (haosdent/vllm's 16-job variant of the allover326 Dockerfile.fullbuild) into the repo as docker/Dockerfile.sm80, so images build from this repo instead of a separately-tracked Dockerfile plus bind-mounted fixes. Adds the safetensors F8_E8M0 overlay COPY step and an in-container import gate for vllm, the custom ops, and the DeepSeek-V4 text/vision/dspark classes, so a broken import fails the build instead of surfacing only at boot. Build from the repo root: `docker build -f docker/Dockerfile.sm80 -t .` --- docker/Dockerfile.sm80 | 54 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 docker/Dockerfile.sm80 diff --git a/docker/Dockerfile.sm80 b/docker/Dockerfile.sm80 new file mode 100644 index 000000000000..953388339c73 --- /dev/null +++ b/docker/Dockerfile.sm80 @@ -0,0 +1,54 @@ +# Full CUDA source build of the `sm80` branch for NVIDIA Ampere +# (compute capability 8.0, e.g. CMP 170HX / A100). Build from the repo +# root: `docker build -f docker/Dockerfile.sm80 -t .` +# +# Lineage: haosdent/vllm@dsv4-flash-a100 Dockerfile.fullbuild16 (16-job +# variant of the allover326 Dockerfile.fullbuild), same torch pin and +# SM80-only arch list. See docs/SM80.md for branch provenance. +FROM nvidia/cuda:13.0.2-cudnn-devel-ubuntu24.04 + +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y --no-install-recommends \ + python3.12 python3.12-dev python3.12-venv git build-essential curl \ + ca-certificates ccache cmake ninja-build \ + && rm -rf /var/lib/apt/lists/* +RUN python3.12 -m venv /opt/venv +ENV PATH=/opt/venv/bin:/usr/local/cuda/bin:/usr/lib/ccache:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin +ENV CUDA_HOME=/usr/local/cuda + +WORKDIR /vllm +COPY . /vllm/ + +ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0 +ENV TORCH_CUDA_ARCH_LIST=8.0 +ENV MAX_JOBS=16 +ENV CMAKE_BUILD_TYPE=Release +ENV VLLM_TARGET_DEVICE=cuda + +RUN pip install --no-cache-dir -U pip "setuptools>=77.0.3,<81.0.0" wheel \ + "setuptools-scm>=8.0" "setuptools-rust>=1.9.0" ninja "cmake>=3.26.1" \ + "packaging>=24.2" jinja2 +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --default-toolchain 1.95 --profile minimal +ENV PATH=/root/.cargo/bin:$PATH +RUN pip install --no-cache-dir torch==2.13.0 torchaudio==2.11.0 torchvision==0.28.0 + +RUN rustc --version && pip install --no-cache-dir -e . --no-build-isolation + +# Runtime overlay: F8_E8M0 dtype support for the installed safetensors +# package (see patches/README.md). Applied after the editable install so +# it is not clobbered by any vllm dependency re-resolution above. +COPY patches/safetensors_torch_f8_e8m0.py \ + /opt/venv/lib/python3.12/site-packages/safetensors/torch.py + +# Build + import gate: fail the build if the compiled extension, or the +# DeepSeek-V4 text/vision classes and custom ops, do not import. +RUN nvcc --version | tail -2 \ + && python3 -c "import vllm, torch; print('vllm', vllm.__version__, 'torch', torch.__version__)" \ + && ls -la /vllm/vllm/*.so \ + && python3 -c "import torch, vllm; import vllm._custom_ops as o; print('custom ops import OK')" \ + && python3 -c "\ +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4ForCausalLM; \ +from vllm.models.deepseek_v4.nvidia.vl_model import DeepseekV4ForConditionalGeneration; \ +from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM; \ +print('deepseek_v4 text + vision + dspark import OK')" From d2a6da8e9fc203eb6fd1182b94caf11b3a38f53a Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 07:43:06 +0000 Subject: [PATCH 14/15] docker: scope import gate to a driver-less build environment docker build here runs under the default runc runtime with no NVIDIA driver visible, so the deepseek_v4 import chain pulls in vllm/v1/attention/ops/fp8_sm80.py, whose triton-disabled fallback path calls tl.constexpr unconditionally. That call is not itself callable once triton_utils disables triton for a missing driver, so it raises TypeError at import time -- a pre-existing upstream gap, not one of this branch's patches, that only shows up with no GPU present. Narrow the Dockerfile.sm80 build-time gate to vllm/torch/custom-ops import (what a driver-less container can actually verify), and document the fuller deepseek_v4-class import check to run post-build with a GPU runtime once one is free. --- docker/Dockerfile.sm80 | 20 ++++++++++++-------- docs/SM80.md | 21 +++++++++++++++++++++ 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/docker/Dockerfile.sm80 b/docker/Dockerfile.sm80 index 953388339c73..d1468e30b13e 100644 --- a/docker/Dockerfile.sm80 +++ b/docker/Dockerfile.sm80 @@ -41,14 +41,18 @@ RUN rustc --version && pip install --no-cache-dir -e . --no-build-isolation COPY patches/safetensors_torch_f8_e8m0.py \ /opt/venv/lib/python3.12/site-packages/safetensors/torch.py -# Build + import gate: fail the build if the compiled extension, or the -# DeepSeek-V4 text/vision classes and custom ops, do not import. +# Build gate: fail the build if the compiled extension or custom ops do +# not import. Scoped to what a driver-less build container can check — +# `docker build` runs under the default `runc` runtime here, with no +# NVIDIA driver visible, and vllm.models.deepseek_v4 transitively imports +# vllm/v1/attention/ops/fp8_sm80.py, whose triton-disabled fallback path +# calls `tl.constexpr` unconditionally. That call is not itself callable +# once triton_utils disables triton for a missing driver, so it raises +# TypeError at import time — a pre-existing upstream gap (not one of this +# branch's patches) that only shows up with no GPU present. See +# docs/SM80.md for the follow-up import-gate command to run with a GPU +# runtime (`docker run --runtime=nvidia`) once one is free. RUN nvcc --version | tail -2 \ && python3 -c "import vllm, torch; print('vllm', vllm.__version__, 'torch', torch.__version__)" \ && ls -la /vllm/vllm/*.so \ - && python3 -c "import torch, vllm; import vllm._custom_ops as o; print('custom ops import OK')" \ - && python3 -c "\ -from vllm.models.deepseek_v4.nvidia.model import DeepseekV4ForCausalLM; \ -from vllm.models.deepseek_v4.nvidia.vl_model import DeepseekV4ForConditionalGeneration; \ -from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM; \ -print('deepseek_v4 text + vision + dspark import OK')" + && python3 -c "import torch, vllm; import vllm._custom_ops as o; print('custom ops import OK')" diff --git a/docs/SM80.md b/docs/SM80.md index 99f61e3668ec..c650885ca608 100644 --- a/docs/SM80.md +++ b/docs/SM80.md @@ -41,6 +41,27 @@ it exists to carry Ampere fallback kernels that upstream does not (yet) ship. `docker/Dockerfile.sm80`; images no longer need bind-mounted fixes on top of a pinned base. +### Build-time import gate scope + +`docker/Dockerfile.sm80`'s final `RUN` step checks that `vllm`, `torch`, +and the compiled custom ops import. It does not import +`vllm.models.deepseek_v4` classes, because `docker build` runs under +the `runc` runtime with no NVIDIA driver visible, and that import chain +pulls in `vllm/v1/attention/ops/fp8_sm80.py`. That module's +triton-disabled fallback path calls `tl.constexpr` unconditionally, +which is not callable once `triton_utils` disables triton for a +missing driver — a pre-existing upstream gap, not one of this branch's +patches, that only shows up with no GPU present. Run the fuller check +after the build, with a GPU runtime, once one is free: + +``` +docker run --rm --runtime=nvidia python3 -c " +from vllm.models.deepseek_v4.nvidia.model import DeepseekV4ForCausalLM +from vllm.models.deepseek_v4.nvidia.vl_model import DeepseekV4ForConditionalGeneration +from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM +print('deepseek_v4 text + vision + dspark import OK')" +``` + ### Known limits (vision build) - **`c>=4` crash**: an EngineCore/`shm_broadcast` timeout at concurrency From 7060f0009e071da545def980ca305c8954153cb2 Mon Sep 17 00:00:00 2001 From: Sean Phan Date: Thu, 3 Sep 2026 09:11:48 +0000 Subject: [PATCH 15/15] docker: exclude non-compile paths from build context, document why COPY . /vllm/ sits before the ~80-minute compile RUN in docker/Dockerfile.sm80, so any change inside the build context -- the Dockerfile itself, a doc edit, even a new commit touching .git -- busts that cache and forces a full recompile. Add docker/, notebooks/, *.md, and receipts/ to .dockerignore so edits there stop invalidating the compile layer. Keep .git in the context; setuptools-scm reads it for the version string. Document the cause and a follow-up (split the gate/overlay steps into a later stage with a narrow COPY) in docs/SM80.md. --- .dockerignore | 6 +++++- docs/SM80.md | 19 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/.dockerignore b/.dockerignore index ce6de41a2cd3..8cf9f366bfcd 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,3 +1,7 @@ -.git __pycache__ *.pyc +docker/ +notebooks/ +*.md +receipts/ +**/receipts/ diff --git a/docs/SM80.md b/docs/SM80.md index c650885ca608..922bc1e07dde 100644 --- a/docs/SM80.md +++ b/docs/SM80.md @@ -62,6 +62,25 @@ from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM print('deepseek_v4 text + vision + dspark import OK')" ``` +### Build cache + +`docker/Dockerfile.sm80` runs `COPY . /vllm/` before the compile step +(`pip install -e .`), which takes about 80 minutes on SM80 hardware. Any +change inside the build context invalidates that layer and forces a full +recompile — this includes edits to the Dockerfile itself, to unrelated +docs, or to any tracked file (a new commit changes `.git`, which +`setuptools-scm` reads for the version string). + +A repo-root `.dockerignore` now excludes `docker/`, `notebooks/`, `*.md`, +and `receipts/` from the build context, so edits to those paths do not +bust the compile cache. `.git` stays in the context on purpose — +`setuptools-scm` needs it to derive the package version. + +For a future revision of this Dockerfile, move the import-gate and +safetensors-overlay steps into a separate, later build stage with a +narrow `COPY` (just the gate script and the patch file), so gate-only +changes never touch the layer that triggers the compile. + ### Known limits (vision build) - **`c>=4` crash**: an EngineCore/`shm_broadcast` timeout at concurrency