diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py index 46c887838a..1418704036 100644 --- a/tests/models/test_registry.py +++ b/tests/models/test_registry.py @@ -48,9 +48,11 @@ def test_registry_imports(model_arch): "(see #41376)" ) - # DSpark draft model is NVIDIA-only; class is stubbed to None on ROCm/XPU. - if model_arch == "DSparkDraftModel" and not current_platform.is_cuda(): - pytest.skip("DSparkDraftModel is only supported on CUDA") + # DSpark draft model is supported on CUDA and ROCm; stubbed to None on XPU. + if model_arch == "DSparkDraftModel" and not ( + current_platform.is_cuda() or current_platform.is_rocm() + ): + pytest.skip("DSparkDraftModel is only supported on CUDA and ROCm") # Ensure all model classes can be imported successfully model_cls = ModelRegistry._try_load_model_cls(model_arch) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index ea749091da..b9c6126219 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -942,6 +942,31 @@ class SpeculativeConfig: "`num_speculative_tokens` was not provided" ) + if self.method == "dspark": + # DSpark is a semi-autoregressive *block* drafter. A + # speculative length smaller than the checkpoint's block + # feeds the block / Markov-head machinery an unsupported + # layout and yields incorrect (garbled) output rather than + # merely lower acceptance. Require num_speculative_tokens to + # be at least the block size (e.g. 5 or 7 for DeepSeek-V4). + dspark_block_size = getattr( + self.draft_model_config.hf_config, + "dspark_block_size", + None, + ) + if ( + dspark_block_size is not None + and self.num_speculative_tokens < dspark_block_size + ): + raise ValueError( + "DSpark requires num_speculative_tokens >= " + f"dspark_block_size ({dspark_block_size}); got " + f"{self.num_speculative_tokens}. Smaller values " + "produce incorrect output. Use " + f"num_speculative_tokens={dspark_block_size} or " + "larger (e.g. 7)." + ) + self.draft_tensor_parallel_size = ( SpeculativeConfig._verify_and_get_draft_tp( self.target_parallel_config, diff --git a/vllm/models/deepseek_v4/__init__.py b/vllm/models/deepseek_v4/__init__.py index d03eba3d41..05662bbd32 100644 --- a/vllm/models/deepseek_v4/__init__.py +++ b/vllm/models/deepseek_v4/__init__.py @@ -15,16 +15,16 @@ from .quant_config import DeepseekV4FP8Config # default that mypy sees; the ROCm/XPU branches override at runtime and are # kept type-compatible via ``# type: ignore[assignment]``. if current_platform.is_rocm(): + from .amd.dspark import ( # type: ignore[assignment] + DSparkDeepseekV4ForCausalLM, + ) from .amd.model import DeepseekV4ForCausalLM from .amd.mtp import DeepSeekV4MTP - - # DSpark is NVIDIA-only for now. - DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment] elif current_platform.is_xpu(): from .xpu.model import DeepseekV4ForCausalLM # type: ignore[assignment] from .xpu.mtp import DeepSeekV4MTP # type: ignore[assignment] - DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment] + DSparkDeepseekV4ForCausalLM = None # type: ignore[assignment, misc] else: from .nvidia.dspark import ( # type: ignore[assignment] DSparkDeepseekV4ForCausalLM, diff --git a/vllm/models/deepseek_v4/amd/dspark.py b/vllm/models/deepseek_v4/amd/dspark.py new file mode 100644 index 0000000000..4769483ff2 --- /dev/null +++ b/vllm/models/deepseek_v4/amd/dspark.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DSpark draft model for DeepSeek-V4 on ROCm/AMD (gfx950). + +ROCm port of ``nvidia/dspark.py``. Follows the same nvidia->amd recipe used for +``amd/mtp.py``: + + * import ``DeepseekV4DecoderLayer`` from the AMD ``.model`` (aiter/triton + attention + MHC CustomOp path) instead of the nvidia one; + * route the MHC head through the ``HCHeadOp`` CustomOp dispatcher (aiter / + tilelang / triton / torch) instead of calling the tilelang kernels directly, + and gate the trailing ``mhc_post`` on ``use_fused_mhc`` (False on the aiter + path, where the decoder layer already applies hc_post in-layer); + * drop the mega-MoE weight path (``make_deepseek_v4_expert_params_mapping`` / + ``use_mega_moe`` / ``finalize_mega_moe_weights`` do not exist in amd/model.py). + +Everything else — the semi-autoregressive drafting hooks, the Markov head, the +sliding-window context-KV insert, and the checkpoint ``mtp.*`` weight remap — is +pure torch / Triton and shared with the nvidia implementation unchanged. +""" + +from collections.abc import Iterable + +import regex as re +import torch +import torch.nn as nn + +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.mhc import HCHeadOp +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3_dspark import ( + DSparkMarkovHead, +) +from vllm.model_executor.models.utils import maybe_prefix + +from .model import ( + DeepseekV4DecoderLayer, +) + +logger = init_logger(__name__) + +# MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders): +# fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``. +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +class DSparkDeepseekV4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + config = vllm_config.speculative_config.draft_model_config.hf_config + self.config = config + self.hidden_size = config.hidden_size + self.hc_mult = config.hc_mult + self.hc_eps = config.hc_eps + self.rms_norm_eps = config.rms_norm_eps + self.num_hidden_layers = config.num_hidden_layers + self.target_layer_ids = tuple(config.dspark_target_layer_ids) + + self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 + + # Shared with the target (aliased by the speculator's loading utility). + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "embed_tokens"), + ) + + self.main_proj = ReplicatedLinear( + config.hidden_size * len(self.target_layer_ids), + config.hidden_size, + bias=False, + return_bias=False, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "main_proj"), + ) + self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + current_vllm_config = get_current_vllm_config() + self.layers = nn.ModuleList( + [ + DeepseekV4DecoderLayer( + current_vllm_config, + prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), + ) + for i in range(self.num_dspark_layers) + ] + ) + + # Heads: final norm + hc_head, and the Markov head + # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + draft_vocab_size = ( + getattr(config, "draft_vocab_size", None) or config.vocab_size + ) + self.markov_head = DSparkMarkovHead( + config.vocab_size, + draft_vocab_size, + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + + # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch), + # replacing the direct nvidia tilelang kernel call. + self.hc_head_op = HCHeadOp() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + """main_x = main_norm(main_proj(concat of target aux hidden states)). + + ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)]. + """ + return self.main_norm(self.main_proj(aux_hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + main_x: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + """Insert the sliding-window context KV for every draft layer. + + Mirrors the reference DSparkAttention: each layer derives its context KV + from the SAME projected target hidden ``main_x``, via that layer's own + ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the + layer's context slots. + + ``context_slot_mappings`` is a per-layer list (each entry is the context + slot mapping for that layer's kv-cache group, since the hybrid manager may + place draft layers in different groups). ``None`` (or a ``None`` entry) + runs the projection to reserve workspace but writes nothing (profiling). + """ + for i, layer in enumerate(self.layers): + slot_mapping = ( + None if context_slot_mappings is None else context_slot_mappings[i] + ) + attn = layer.attn + # Optimized DSV4 MLA path: wkv part of the fused wq_a|wkv projection + # (q_lora part discarded), then RoPE/quant/insert via the fused op. + qr_kv, _ = attn.fused_wqa_wkv(main_x) + kv = qr_kv[..., attn.q_lora_rank :] + kv = attn.kv_norm(kv) + if slot_mapping is None: + continue + _insert_context_kv(attn, kv, context_positions, slot_mapping) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). + hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) + + residual = post_mix = res_mix = None + for layer in self.layers: + hidden_states, residual, post_mix, res_mix = layer( + hidden_states, + positions, + input_ids, + post_mix, + res_mix, + residual, + ) + # On the fused-MHC path the trailing hc_post must be applied here; on the + # aiter unfused path (ROCm default) the decoder layer already applied + # hc_post in-layer and returned None mixes, so this is skipped. Mirrors + # amd/mtp.py. + last_layer = self.layers[-1] + if last_layer.use_fused_mhc: + hidden_states = last_layer.hc_post( + hidden_states, residual, post_mix, res_mix + ) + # hc_head reduces the hc copies; return the PRE-norm head hidden. + hidden_states = self.hc_head_op( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + return hidden_states + + +def _insert_context_kv( + attn: nn.Module, + kv: torch.Tensor, + positions: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. + + Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy + query and discard it, since context tokens have no query). Mirrors + ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``. + """ + swa_cache = attn.swa_cache_layer.kv_cache + block_size = attn.swa_cache_layer.block_size + cos_sin_cache = attn.rotary_emb.cos_sin_cache + cache_dtype = swa_cache.dtype + n_ctx = kv.shape[0] + dummy_q = torch.zeros( + (n_ctx, attn.n_local_heads, attn.head_dim), + dtype=kv.dtype, + device=kv.device, + ) + if cache_dtype == torch.uint8: + # fp8_ds_mla UE8M0 paged layout (the gfx950 aiter SWA cache path). + swa_2d = swa_cache.view(swa_cache.shape[0], -1) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + dummy_q, + kv, + swa_2d, + slot_mapping, + positions, + cos_sin_cache, + attn.padded_heads, + attn.eps, + block_size, + ) + elif cache_dtype == torch.bfloat16: + swa_3d = swa_cache.view(-1, block_size, attn.head_dim) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + dummy_q, + kv, + swa_3d, + slot_mapping, + positions, + cos_sin_cache, + attn.eps, + block_size, + ) + else: # per-tensor fp8 (torch.float8_e4m3fn) + # NOTE(rocm): unreachable on ROCm/aiter, where the SWA cache dtype is + # uint8 (fp8_ds_mla) or bfloat16. This branch relies on FlashInfer-only + # attributes (``_flashinfer_fp8_*``) that the aiter attention layer does + # not define; kept for parity with the nvidia path. + swa_3d = swa_cache.view(-1, block_size, attn.head_dim) + dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + dummy_q, + kv, + dummy_q_fp8, + swa_3d, + slot_mapping, + positions, + cos_sin_cache, + attn._flashinfer_fp8_kv_scale, + attn._flashinfer_fp8_q_scale_inv, + attn.eps, + block_size, + ) + + +class DSparkDeepseekV4ForCausalLM(nn.Module): + # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so + # load_dspark_model always aliases the target's. + has_own_embed_tokens = False + has_own_lm_head = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + self.model = DSparkDeepseekV4Model( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + # Shared with the target (aliased by the speculator's load utility). + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(self.config.vocab_size) + + # --- Hooks used by the speculator ------------------------------------- + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(aux_hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + # DSV4 MLA path: each draft layer's sliding-window cache is a separate + # layer, named by its prefix. + return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mappings: list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mappings + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + # Returns the pre-norm hc_head hidden ([T, hidden_size]). + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Base logits U_k = lm_head(norm(head_hidden)).""" + return self.logits_processor(self.lm_head, self.model.norm(hidden_states)) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + # Full-vocab draft: base logits, no d2t scatter. + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_ids # full-vocab: draft ids are target ids + + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.bias(markov_embed, self.logits_processor) + + # --- Weight loading ---------------------------------------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint. + + Non-mtp weights (embed/head/main layers) belong to the target model and + are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. + """ + # AMD DeepseekV4MoE has no mega-MoE path; always use the standard + # per-expert fused-MoE mapping (mirrors amd/mtp.py). + expert_mapping = fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + expert_scale_suffix = ( + ".weight_scale" + if getattr(self.config, "expert_dtype", "fp4") == "fp4" + else ".weight_scale_inv" + ) + + # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params. + stacked_params_mapping = [ + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_local_head = self.config.num_attention_heads // tp_size + head_start = n_local_head * tp_rank + head_end = n_local_head * (tp_rank + 1) + + for name, loaded_weight in weights: + mapped = self._remap_dspark_name(name) + if mapped is None: + continue + name = mapped + + # ``.scale`` -> per-method scale suffix. + if name.endswith(".scale"): + suffix = ( + expert_scale_suffix + if _EXPERT_SCALE_RE.search(name) + else ".weight_scale_inv" + ) + name = name.removesuffix(".scale") + suffix + + # E8M0 expert scales: keep raw exponent bytes. + if ".experts." in name: + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + loaded_weight = loaded_weight.view(torch.uint8) + for param_name, weight_name, expert_id, shard_id in expert_mapping: + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + param = params_dict[name_mapped] + success = param.weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(name_mapped) + break + continue + + # Stacked rules only apply to decoder-layer weights. Head-stack params + # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g. + # "markov_w1" would collide with the "w1" shard rule. + is_layer_param = name.startswith("model.layers.") + for param_name, weight_name, stacked_shard_id in stacked_params_mapping: + if not is_layer_param or weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + param.weight_loader(param, loaded_weight, stacked_shard_id) + loaded_params.add(name) + break + else: + if "attn_sink" in name: + narrow = loaded_weight[head_start:head_end] + params_dict[name][: narrow.shape[0]].copy_(narrow) + loaded_params.add(name) + continue + if ".shared_experts.w2" in name: + name = name.replace( + ".shared_experts.w2", ".shared_experts.down_proj" + ) + if name.endswith(".ffn.gate.bias"): + name = name.replace( + ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" + ) + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) + return loaded_params + + def _remap_dspark_name(self, name: str) -> str | None: + """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. + + Returns None for non-mtp weights (owned by the target model). + """ + m = re.match(r"mtp\.(\d+)\.(.*)", name) + if m is None: + return None + stage = int(m.group(1)) + rest = m.group(2) + # The confidence head is not wired into inference yet; drop its weights. + if rest.startswith("confidence_head."): + return None + # Head-stack params live at model level (mtp.last), context combiner at + # model level (mtp.0); everything else is a per-layer decoder block. + head_prefixes = ( + "norm.", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + "markov_head.", + ) + if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( + head_prefixes + ): + return f"model.{rest}" + return f"model.layers.{stage}.{rest}" diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 01646ad994..9093bd4128 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -40,7 +40,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + SupportsEagle3, + SupportsPP, +) from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, @@ -437,7 +441,7 @@ class DeepseekV4DecoderLayer(nn.Module): ) -class DeepseekV4Model(nn.Module): +class DeepseekV4Model(nn.Module, EagleModelMixin): def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -573,7 +577,19 @@ class DeepseekV4Model(nn.Module): hidden_states = intermediate_tensors["hidden_states"] residual, post_mix, res_mix = None, None, None - for layer in islice(self.layers, self.start_layer, self.end_layer): + # EAGLE3 / DSpark / DFlash aux hidden states: reconstructed (post-mhc) + # hidden state at the configured target layers, averaged over the + # hc_mult streams to [T, hidden_size]. Empty unless a draft model set + # aux_hidden_state_layers. + aux_hidden_states: list[torch.Tensor] = [] + # On the fused path the final layer's hc_post output is reused below + # (avoids computing hc_post twice when the last layer is also an aux + # layer). + final_aux_recon: torch.Tensor | None = None + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer), + start=self.start_layer, + ): hidden_states, residual, post_mix, res_mix = layer( hidden_states, positions, @@ -582,8 +598,30 @@ class DeepseekV4Model(nn.Module): res_mix, residual, ) + if (idx + 1) in self.aux_hidden_state_layers: + # On the unfused (aiter) path the layer already applied hc_post, + # so hidden_states is the reconstructed stream; on the fused + # path reconstruct it via hc_post before averaging. + if layer.use_fused_mhc: + aux_recon = layer.hc_post( + hidden_states, residual, post_mix, res_mix + ) + final_aux_recon = aux_recon + else: + aux_recon = hidden_states + aux_hidden_states.append(aux_recon.mean(dim=1)) if layer is not None and layer.use_fused_mhc: - hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) + # Reuse the last layer's hc_post output if it was already computed + # for the aux hidden state above; otherwise compute it now. + if ( + final_aux_recon is not None + and self.end_layer in self.aux_hidden_state_layers + ): + hidden_states = final_aux_recon + else: + hidden_states = layer.hc_post( + hidden_states, residual, post_mix, res_mix + ) if not get_pp_group().is_last_rank: return IntermediateTensors({"hidden_states": hidden_states}) @@ -601,6 +639,8 @@ class DeepseekV4Model(nn.Module): self.hc_eps, ) hidden_states = self.norm(hidden_states) + if len(aux_hidden_states) > 0: + return hidden_states, aux_hidden_states return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -751,7 +791,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: ) -class DeepseekV4ForCausalLM(nn.Module, SupportsPP): +class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3): model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. diff --git a/vllm/models/deepseek_v4/amd/rocm.py b/vllm/models/deepseek_v4/amd/rocm.py index b3456ba785..d4aef6ec9e 100644 --- a/vllm/models/deepseek_v4/amd/rocm.py +++ b/vllm/models/deepseek_v4/amd/rocm.py @@ -371,8 +371,12 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens + # The non-causal (DSpark draft) path widens each token's SWA index list + # to ``noncausal_index_width`` (>= window_size), so size the persistent + # ragged buffer to the wider bound to cover both causal and non-causal. + swa_index_width = max(self.window_size, self.noncausal_index_width) self.decode_swa_ragged_indices_buffer = torch.empty( - max_tokens * self.window_size, + max_tokens * swa_index_width, dtype=torch.int32, device=self.device, ) @@ -411,7 +415,9 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuild self.decode_swa_ragged_indices_buffer, self.decode_swa_ragged_indptr_buffer, base.num_decode_tokens, - self.window_size, + # Actual dense width for this build: window_size (causal) or + # noncausal_index_width (DSpark non-causal draft). + base.decode_swa_indices.shape[-1], ) return DeepseekV4ROCMAiterSparseSWAMetadata(