[ROCm] Enable DeepSeek-V4 DSpark speculative decoding on AMD (MI350X / MI355X, gfx950) (#47419)
Signed-off-by: larryli2-amd <larryli2@amd.com> Signed-off-by: larryli2-amd <Larry.Li@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
larryli2-amd
co-authored by
Andreas Karatzas
Andreas Karatzas
parent
08dfd68610
commit
c227aaa3f8
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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}"
|
||||
@@ -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.
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user