| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DSpark draft model for DeepSeek-V4 (semi-autoregressive speculative decoding). |
| 4 | |
| 5 | See: qwen3_dspark.py for base architecture. This one is specialized to the DSV4 DSpark, |
| 6 | which reuses the target model's architecture similarly to MTP. |
| 7 | |
| 8 | To implement non-causal attention, we leverage the sparse attention implementation to |
| 9 | include the future query tokens in the top-k indices for each query token. |
| 10 | """ |
| 11 | |
| 12 | from collections.abc import Iterable |
| 13 | |
| 14 | import regex as re |
| 15 | import torch |
| 16 | import torch.nn as nn |
| 17 | |
| 18 | import vllm.envs as envs |
| 19 | from vllm.config import VllmConfig, get_current_vllm_config |
| 20 | from vllm.config.kernel import MEGA_MOE_BACKENDS |
| 21 | from vllm.distributed import ( |
| 22 | get_tensor_model_parallel_rank, |
| 23 | get_tensor_model_parallel_world_size, |
| 24 | ) |
| 25 | from vllm.forward_context import get_forward_context, is_forward_context_available |
| 26 | from vllm.logger import init_logger |
| 27 | from vllm.model_executor.kernels.mhc.tilelang import ( |
| 28 | hc_head_fused_kernel_tilelang, |
| 29 | mhc_post_tilelang, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe import ( |
| 32 | fused_moe_make_expert_params_mapping, |
| 33 | ) |
| 34 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 35 | from vllm.model_executor.layers.linear import ( |
| 36 | MergedColumnParallelLinear, |
| 37 | ReplicatedLinear, |
| 38 | ) |
| 39 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 40 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 41 | ParallelLMHead, |
| 42 | VocabParallelEmbedding, |
| 43 | ) |
| 44 | from vllm.model_executor.model_loader.weight_utils import default_weight_loader |
| 45 | from vllm.model_executor.models.qwen3_dspark import ( |
| 46 | DSparkConfidenceHead, |
| 47 | DSparkMarkovHead, |
| 48 | ) |
| 49 | from vllm.model_executor.models.utils import maybe_prefix |
| 50 | from vllm.models.common.ops.sequence_parallel import ( |
| 51 | sp_all_gather, |
| 52 | sp_padding_mask, |
| 53 | sp_shard, |
| 54 | ) |
| 55 | from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID |
| 56 | |
| 57 | from .model import ( |
| 58 | DeepseekV4DecoderLayer, |
| 59 | DeepseekV4Model, |
| 60 | _use_sequence_parallel, |
| 61 | make_deepseek_v4_expert_params_mapping, |
| 62 | prepare_mega_gate_routing_metadata, |
| 63 | ) |
| 64 | |
| 65 | logger = init_logger(__name__) |
| 66 | |
| 67 | # MoE expert scale suffix differs by expert dtype (mirrors deepseek_v4 loaders): |
| 68 | # fp4 experts register ``.weight_scale``; block-fp8 experts ``.weight_scale_inv``. |
| 69 | _EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") |
| 70 | _CONTEXT_WKV_RE = re.compile(r"^mtp\.(\d+)\.attn\.wkv\.(.+)$") |
| 71 | |
| 72 | |
| 73 | def _duplicate_context_wkv_weights( |
| 74 | weights: Iterable[tuple[str, torch.Tensor]], num_layers: int |
| 75 | ) -> Iterable[tuple[str, torch.Tensor]]: |
| 76 | """Load every draft layer's WKV into the cross-layer projection.""" |
| 77 | for name, weight in weights: |
| 78 | yield name, weight |
| 79 | match = _CONTEXT_WKV_RE.fullmatch(name) |
| 80 | if match is None: |
| 81 | continue |
| 82 | layer_idx = int(match.group(1)) |
| 83 | if layer_idx >= num_layers: |
| 84 | continue |
| 85 | stacked_weight = weight.detach() |
| 86 | stacked_weight.shard_id = layer_idx |
| 87 | yield f"context_wkv_proj.{match.group(2)}", stacked_weight |
| 88 | |
| 89 | |
| 90 | class DSparkDeepseekV4Model(nn.Module): |
| 91 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 92 | super().__init__() |
| 93 | assert vllm_config.speculative_config is not None |
| 94 | config = vllm_config.speculative_config.draft_model_config.hf_config |
| 95 | self.config = config |
| 96 | self.hidden_size = config.hidden_size |
| 97 | self.hc_mult = config.hc_mult |
| 98 | self.hc_eps = config.hc_eps |
| 99 | self.rms_norm_eps = config.rms_norm_eps |
| 100 | self.num_hidden_layers = config.num_hidden_layers |
| 101 | self.target_layer_ids = tuple(config.dspark_target_layer_ids) |
| 102 | self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS |
| 103 | self.use_sequence_parallel = _use_sequence_parallel(vllm_config) |
| 104 | |
| 105 | self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3 |
| 106 | |
| 107 | # Shared with the target (aliased by the speculator's loading utility). |
| 108 | self.embed_tokens = VocabParallelEmbedding( |
| 109 | config.vocab_size, |
| 110 | config.hidden_size, |
| 111 | prefix=maybe_prefix(prefix, "embed_tokens"), |
| 112 | ) |
| 113 | |
| 114 | self.main_proj = ReplicatedLinear( |
| 115 | config.hidden_size * len(self.target_layer_ids), |
| 116 | config.hidden_size, |
| 117 | bias=False, |
| 118 | return_bias=False, |
| 119 | quant_config=vllm_config.quant_config, |
| 120 | prefix=maybe_prefix(prefix, "main_proj"), |
| 121 | ) |
| 122 | self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 123 | |
| 124 | self.topk_indices_buffer = torch.empty( |
| 125 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 126 | config.index_topk, |
| 127 | dtype=torch.int32, |
| 128 | ) |
| 129 | |
| 130 | current_vllm_config = get_current_vllm_config() |
| 131 | self.layers = nn.ModuleList( |
| 132 | [ |
| 133 | DeepseekV4DecoderLayer( |
| 134 | current_vllm_config, |
| 135 | prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), |
| 136 | topk_indices_buffer=self.topk_indices_buffer, |
| 137 | ) |
| 138 | for i in range(self.num_dspark_layers) |
| 139 | ] |
| 140 | ) |
| 141 | self.context_wkv_proj = MergedColumnParallelLinear( |
| 142 | config.hidden_size, |
| 143 | [config.head_dim] * self.num_dspark_layers, |
| 144 | bias=False, |
| 145 | return_bias=False, |
| 146 | quant_config=vllm_config.quant_config, |
| 147 | prefix=maybe_prefix(prefix, "context_wkv_proj"), |
| 148 | disable_tp=True, |
| 149 | ) |
| 150 | |
| 151 | # Heads: final norm + hc_head, and the Markov + confidence heads |
| 152 | # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint |
| 153 | self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) |
| 154 | hc_dim = self.hc_mult * config.hidden_size |
| 155 | self.hc_head_fn = nn.Parameter( |
| 156 | torch.empty(self.hc_mult, hc_dim, dtype=torch.float32), |
| 157 | requires_grad=False, |
| 158 | ) |
| 159 | self.hc_head_base = nn.Parameter( |
| 160 | torch.empty(self.hc_mult, dtype=torch.float32), requires_grad=False |
| 161 | ) |
| 162 | self.hc_head_scale = nn.Parameter( |
| 163 | torch.empty(1, dtype=torch.float32), requires_grad=False |
| 164 | ) |
| 165 | draft_vocab_size = ( |
| 166 | getattr(config, "draft_vocab_size", None) or config.vocab_size |
| 167 | ) |
| 168 | self.markov_head = DSparkMarkovHead( |
| 169 | config.vocab_size, |
| 170 | draft_vocab_size, |
| 171 | config.dspark_markov_rank, |
| 172 | prefix=maybe_prefix(prefix, "markov_head"), |
| 173 | ) |
| 174 | self.confidence_head: DSparkConfidenceHead | None = None |
| 175 | if getattr(config, "enable_confidence_head", True): |
| 176 | self.confidence_head = DSparkConfidenceHead( |
| 177 | config.hidden_size + config.dspark_markov_rank, |
| 178 | prefix=maybe_prefix(prefix, "confidence_head"), |
| 179 | ) |
| 180 | |
| 181 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 182 | return self.embed_tokens(input_ids) |
| 183 | |
| 184 | def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: |
| 185 | """main_x = main_norm(main_proj(concat of target aux hidden states)). |
| 186 | |
| 187 | ``aux_hidden_states`` is [T, hidden_size * len(target_layer_ids)]. |
| 188 | """ |
| 189 | return self.main_norm(self.main_proj(aux_hidden_states)) |
| 190 | |
| 191 | @torch.inference_mode() |
| 192 | def precompute_and_store_context_kv( |
| 193 | self, |
| 194 | main_x: torch.Tensor, |
| 195 | context_positions: torch.Tensor, |
| 196 | context_slot_mappings: list[torch.Tensor | None] | None = None, |
| 197 | ) -> None: |
| 198 | """Insert the sliding-window context KV for every draft layer. |
| 199 | |
| 200 | Mirrors the reference DSparkAttention: each layer derives its context KV |
| 201 | from the SAME projected target hidden ``main_x``, via that layer's own |
| 202 | ``wkv`` + ``kv_norm`` + RoPE + quant, then writes it at the |
| 203 | layer's context slots. |
| 204 | |
| 205 | ``context_slot_mappings`` is a per-layer list (each entry is the context |
| 206 | slot mapping for that layer's kv-cache group, since the hybrid manager may |
| 207 | place draft layers in different groups). ``None`` (or a ``None`` entry) |
| 208 | runs the projection to reserve workspace but writes nothing (profiling). |
| 209 | """ |
| 210 | all_kv = self.context_wkv_proj(main_x).view( |
| 211 | main_x.shape[0], self.num_dspark_layers, self.config.head_dim |
| 212 | ) |
| 213 | for i, (layer, kv) in enumerate( |
| 214 | zip(self.layers, all_kv.unbind(1), strict=True) |
| 215 | ): |
| 216 | slot_mapping = ( |
| 217 | None if context_slot_mappings is None else context_slot_mappings[i] |
| 218 | ) |
| 219 | attn = layer.attn |
| 220 | kv = attn.kv_norm(kv) |
| 221 | if slot_mapping is None: |
| 222 | continue |
| 223 | _insert_context_kv(attn, kv, context_positions, slot_mapping) |
| 224 | |
| 225 | def forward( |
| 226 | self, |
| 227 | input_ids: torch.Tensor, |
| 228 | positions: torch.Tensor, |
| 229 | inputs_embeds: torch.Tensor | None = None, |
| 230 | ) -> torch.Tensor: |
| 231 | if inputs_embeds is None: |
| 232 | inputs_embeds = self.embed_input_ids(input_ids) |
| 233 | full_num_tokens = positions.shape[0] |
| 234 | if self.use_sequence_parallel: |
| 235 | if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): |
| 236 | forward_context = get_forward_context() |
| 237 | forward_context.is_padding = sp_padding_mask( |
| 238 | forward_context.is_padding, inputs_embeds |
| 239 | ) |
| 240 | inputs_embeds = sp_shard(inputs_embeds) |
| 241 | input_ids = sp_shard(input_ids) |
| 242 | mega_gate_metadata = None |
| 243 | if self.use_mega_moe: |
| 244 | mega_gate_metadata = prepare_mega_gate_routing_metadata( |
| 245 | input_ids, |
| 246 | has_hash_routing=False, |
| 247 | image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID |
| 248 | if getattr(self.config, "vision_n_layers", 0) > 0 |
| 249 | else None, |
| 250 | ) |
| 251 | # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]). |
| 252 | hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1) |
| 253 | |
| 254 | residual = post_mix = res_mix = None |
| 255 | for layer in self.layers: |
| 256 | hidden_states, residual, post_mix, res_mix = layer( |
| 257 | hidden_states, |
| 258 | positions, |
| 259 | input_ids, |
| 260 | post_mix, |
| 261 | res_mix, |
| 262 | residual, |
| 263 | mega_gate_metadata=mega_gate_metadata, |
| 264 | ) |
| 265 | hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix) |
| 266 | if self.use_sequence_parallel: |
| 267 | hidden_states = sp_all_gather(hidden_states)[:full_num_tokens] |
| 268 | # hc_head reduces the hc copies; return the PRE-norm head hidden |
| 269 | hidden_states = hc_head_fused_kernel_tilelang( |
| 270 | hidden_states, |
| 271 | self.hc_head_fn, |
| 272 | self.hc_head_scale, |
| 273 | self.hc_head_base, |
| 274 | self.rms_norm_eps, |
| 275 | self.hc_eps, |
| 276 | ) |
| 277 | return hidden_states |
| 278 | |
| 279 | |
| 280 | def _insert_context_kv( |
| 281 | attn: nn.Module, |
| 282 | kv: torch.Tensor, |
| 283 | positions: torch.Tensor, |
| 284 | slot_mapping: torch.Tensor, |
| 285 | ) -> None: |
| 286 | """RoPE + quant + paged-cache insert of (already kv_norm'd) context KV. |
| 287 | |
| 288 | Reuses the DSV4 fused insert ops (which also process a query; we pass a dummy |
| 289 | query and discard it, since context tokens have no query). Mirrors |
| 290 | ``DeepseekV4Attention._fused_qnorm_rope_kv_insert``. |
| 291 | """ |
| 292 | swa_cache = attn.swa_cache_layer.kv_cache |
| 293 | block_size = attn.swa_cache_layer.block_size |
| 294 | cos_sin_cache = attn.rotary_emb.cos_sin_cache |
| 295 | cache_dtype = swa_cache.dtype |
| 296 | n_ctx = kv.shape[0] |
| 297 | dummy_q = torch.zeros( |
| 298 | (n_ctx, attn.n_local_heads, attn.head_dim), |
| 299 | dtype=kv.dtype, |
| 300 | device=kv.device, |
| 301 | ) |
| 302 | if cache_dtype == torch.uint8: |
| 303 | # fp8_ds_mla UE8M0 paged layout |
| 304 | swa_2d = swa_cache.view(swa_cache.shape[0], -1) |
| 305 | torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( |
| 306 | dummy_q, |
| 307 | kv, |
| 308 | swa_2d, |
| 309 | slot_mapping, |
| 310 | positions, |
| 311 | cos_sin_cache, |
| 312 | attn.padded_heads, |
| 313 | attn.eps, |
| 314 | block_size, |
| 315 | ) |
| 316 | elif cache_dtype == torch.bfloat16: |
| 317 | swa_3d = swa_cache.view(-1, block_size, attn.head_dim) |
| 318 | torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( |
| 319 | dummy_q, |
| 320 | kv, |
| 321 | swa_3d, |
| 322 | slot_mapping, |
| 323 | positions, |
| 324 | cos_sin_cache, |
| 325 | attn.eps, |
| 326 | block_size, |
| 327 | ) |
| 328 | else: # per-tensor fp8 (torch.float8_e4m3fn) |
| 329 | # TODO(ben): double-check if this is being dispatched correctly for FI backend |
| 330 | swa_3d = swa_cache.view(-1, block_size, attn.head_dim) |
| 331 | dummy_q_fp8 = torch.zeros_like(dummy_q, dtype=torch.float8_e4m3fn) |
| 332 | torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( |
| 333 | dummy_q, |
| 334 | kv, |
| 335 | dummy_q_fp8, |
| 336 | swa_3d, |
| 337 | slot_mapping, |
| 338 | positions, |
| 339 | cos_sin_cache, |
| 340 | attn._flashinfer_fp8_kv_scale, |
| 341 | attn._flashinfer_fp8_q_scale_inv, |
| 342 | attn.eps, |
| 343 | block_size, |
| 344 | ) |
| 345 | |
| 346 | |
| 347 | class DSparkDeepseekV4ForCausalLM(nn.Module): |
| 348 | # Draft weights ship in the target checkpoint (mtp.*) without embed/head, so |
| 349 | # load_dspark_model always aliases the target's. |
| 350 | has_own_embed_tokens = False |
| 351 | has_own_lm_head = False |
| 352 | # Full-vocab draft: draft ids are target ids, no remapping needed. |
| 353 | draft_id_to_target_id = None |
| 354 | |
| 355 | def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: |
| 356 | super().__init__() |
| 357 | assert vllm_config.speculative_config is not None |
| 358 | self.draft_model_config = vllm_config.speculative_config.draft_model_config |
| 359 | self.config = self.draft_model_config.hf_config |
| 360 | self.quant_config = vllm_config.quant_config |
| 361 | self.pad_shared_expert = getattr( |
| 362 | self.quant_config, "weight_block_size", None |
| 363 | ) is not None and not _use_sequence_parallel(vllm_config) |
| 364 | self.model = DSparkDeepseekV4Model( |
| 365 | vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") |
| 366 | ) |
| 367 | # Shared with the target (aliased by the speculator's load utility). |
| 368 | self.lm_head = ParallelLMHead( |
| 369 | self.config.vocab_size, |
| 370 | self.config.hidden_size, |
| 371 | prefix=maybe_prefix(prefix, "lm_head"), |
| 372 | ) |
| 373 | self.logits_processor = LogitsProcessor(self.config.vocab_size) |
| 374 | |
| 375 | # --- Hooks used by the speculator ------------------------------------- |
| 376 | |
| 377 | def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: |
| 378 | return self.model.embed_input_ids(input_ids) |
| 379 | |
| 380 | def combine_hidden_states(self, aux_hidden_states: torch.Tensor) -> torch.Tensor: |
| 381 | return self.model.combine_hidden_states(aux_hidden_states) |
| 382 | |
| 383 | def get_draft_kv_cache_layer_names(self) -> list[str]: |
| 384 | # DSV4 MLA path: each draft layer's sliding-window cache is a separate |
| 385 | # layer, named by its prefix. |
| 386 | return [layer.attn.swa_cache_layer.prefix for layer in self.model.layers] |
| 387 | |
| 388 | def precompute_and_store_context_kv( |
| 389 | self, |
| 390 | context_states: torch.Tensor, |
| 391 | context_positions: torch.Tensor, |
| 392 | context_slot_mappings: list[torch.Tensor | None] | None = None, |
| 393 | ) -> None: |
| 394 | self.model.precompute_and_store_context_kv( |
| 395 | context_states, context_positions, context_slot_mappings |
| 396 | ) |
| 397 | |
| 398 | def forward( |
| 399 | self, |
| 400 | input_ids: torch.Tensor, |
| 401 | positions: torch.Tensor, |
| 402 | inputs_embeds: torch.Tensor | None = None, |
| 403 | ) -> torch.Tensor: |
| 404 | # Returns the pre-norm hc_head hidden ([T, hidden_size]). |
| 405 | return self.model(input_ids, positions, inputs_embeds) |
| 406 | |
| 407 | def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 408 | """Base logits U_k = lm_head(norm(head_hidden)).""" |
| 409 | return self.logits_processor(self.lm_head, self.model.norm(hidden_states)) |
| 410 | |
| 411 | def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 412 | # Full-vocab draft: base logits, no d2t scatter. |
| 413 | return self.compute_logits(hidden_states) |
| 414 | |
| 415 | def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: |
| 416 | return draft_ids # full-vocab: draft ids are target ids |
| 417 | |
| 418 | def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: |
| 419 | return self.model.markov_head.embed(token_ids) |
| 420 | |
| 421 | def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: |
| 422 | return self.model.markov_head.bias(markov_embed, self.logits_processor) |
| 423 | |
| 424 | def compute_confidence( |
| 425 | self, head_hidden: torch.Tensor, markov_embed: torch.Tensor |
| 426 | ) -> torch.Tensor: |
| 427 | """Per-position acceptance probability for each drafted token.""" |
| 428 | assert self.model.confidence_head is not None |
| 429 | return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed)) |
| 430 | |
| 431 | # --- Weight loading ---------------------------------------------------- |
| 432 | |
| 433 | def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: |
| 434 | """Load the ``mtp.{0,1,2}.*`` draft weights from the target checkpoint. |
| 435 | |
| 436 | Non-mtp weights (embed/head/main layers) belong to the target model and |
| 437 | are skipped here. ``embed_tokens``/``lm_head`` are aliased from the target. |
| 438 | """ |
| 439 | first_layer = self.model.layers[0] |
| 440 | use_mega_moe = first_layer.ffn.use_mega_moe |
| 441 | if use_mega_moe: |
| 442 | expert_mapping = make_deepseek_v4_expert_params_mapping( |
| 443 | self.config.n_routed_experts |
| 444 | ) |
| 445 | else: |
| 446 | expert_mapping = fused_moe_make_expert_params_mapping( |
| 447 | self, |
| 448 | ckpt_gate_proj_name="w1", |
| 449 | ckpt_down_proj_name="w2", |
| 450 | ckpt_up_proj_name="w3", |
| 451 | num_experts=self.config.n_routed_experts, |
| 452 | ) |
| 453 | expert_scale_suffix = ( |
| 454 | ".weight_scale" |
| 455 | if getattr(self.config, "expert_dtype", "fp4") == "fp4" |
| 456 | else ".weight_scale_inv" |
| 457 | ) |
| 458 | |
| 459 | # (param_name, ckpt_shard_name, shard_id) for non-expert stacked params. |
| 460 | stacked_params_mapping = [ |
| 461 | ("gate_up_proj", "w1", 0), |
| 462 | ("gate_up_proj", "w3", 1), |
| 463 | ("attn.fused_wqa_wkv", "attn.wq_a", 0), |
| 464 | ("attn.fused_wqa_wkv", "attn.wkv", 1), |
| 465 | ] |
| 466 | |
| 467 | params_dict = dict(self.named_parameters()) |
| 468 | loaded_params: set[str] = set() |
| 469 | loaded_confidence_head = False |
| 470 | |
| 471 | tp_size = get_tensor_model_parallel_world_size() |
| 472 | tp_rank = get_tensor_model_parallel_rank() |
| 473 | n_local_head = self.config.num_attention_heads // tp_size |
| 474 | head_start = n_local_head * tp_rank |
| 475 | head_end = n_local_head * (tp_rank + 1) |
| 476 | |
| 477 | weights = _duplicate_context_wkv_weights(weights, len(self.model.layers)) |
| 478 | for name, loaded_weight in weights: |
| 479 | mapped = self._remap_dspark_name(name) |
| 480 | if mapped is None: |
| 481 | continue |
| 482 | name = mapped |
| 483 | if "confidence_head." in name: |
| 484 | loaded_confidence_head = True |
| 485 | |
| 486 | # ``.scale`` -> per-method scale suffix. |
| 487 | if name.endswith(".scale"): |
| 488 | suffix = ( |
| 489 | expert_scale_suffix |
| 490 | if _EXPERT_SCALE_RE.search(name) |
| 491 | else ".weight_scale_inv" |
| 492 | ) |
| 493 | name = name.removesuffix(".scale") + suffix |
| 494 | if ".shared_experts.w2" in name: |
| 495 | name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") |
| 496 | if self.pad_shared_expert and ".shared_experts." in name: |
| 497 | loaded_weight = DeepseekV4Model._pad_shared_expert_weight( |
| 498 | self.quant_config, name, loaded_weight |
| 499 | ) |
| 500 | |
| 501 | if name.startswith("model.context_wkv_proj."): |
| 502 | param = params_dict[name] |
| 503 | param.weight_loader(param, loaded_weight, loaded_weight.shard_id) |
| 504 | loaded_params.add(name) |
| 505 | continue |
| 506 | |
| 507 | # E8M0 expert scales: keep raw exponent bytes. |
| 508 | if ".experts." in name: |
| 509 | if ( |
| 510 | "weight_scale" in name |
| 511 | and loaded_weight.dtype == torch.float8_e8m0fnu |
| 512 | ): |
| 513 | loaded_weight = loaded_weight.view(torch.uint8) |
| 514 | for param_name, weight_name, expert_id, shard_id in expert_mapping: |
| 515 | if weight_name not in name: |
| 516 | continue |
| 517 | name_mapped = name.replace(weight_name, param_name) |
| 518 | param = params_dict[name_mapped] |
| 519 | success = param.weight_loader( |
| 520 | param, |
| 521 | loaded_weight, |
| 522 | name_mapped, |
| 523 | shard_id=shard_id, |
| 524 | expert_id=expert_id, |
| 525 | return_success=True, |
| 526 | ) |
| 527 | if success: |
| 528 | loaded_params.add(name_mapped) |
| 529 | break |
| 530 | continue |
| 531 | |
| 532 | # Stacked rules only apply to decoder-layer weights. Head-stack params |
| 533 | # (main_proj/norm/hc_head/markov_head/confidence_head) load directly — |
| 534 | # otherwise e.g. "markov_w1" would collide with the "w1" shard rule. |
| 535 | is_layer_param = name.startswith("model.layers.") |
| 536 | for param_name, weight_name, stacked_shard_id in stacked_params_mapping: |
| 537 | if not is_layer_param or weight_name not in name: |
| 538 | continue |
| 539 | name = name.replace(weight_name, param_name) |
| 540 | param = params_dict[name] |
| 541 | param.weight_loader(param, loaded_weight, stacked_shard_id) |
| 542 | loaded_params.add(name) |
| 543 | break |
| 544 | else: |
| 545 | if "attn_sink" in name: |
| 546 | narrow = loaded_weight[head_start:head_end] |
| 547 | params_dict[name][: narrow.shape[0]].copy_(narrow) |
| 548 | loaded_params.add(name) |
| 549 | continue |
| 550 | if name.endswith(".ffn.gate.bias"): |
| 551 | name = name.replace( |
| 552 | ".ffn.gate.bias", ".ffn.gate.e_score_correction_bias" |
| 553 | ) |
| 554 | param = params_dict[name] |
| 555 | weight_loader = getattr(param, "weight_loader", default_weight_loader) |
| 556 | weight_loader(param, loaded_weight) |
| 557 | loaded_params.add(name) |
| 558 | |
| 559 | if self.model.confidence_head is not None and not loaded_confidence_head: |
| 560 | self.model.confidence_head = None |
| 561 | self.process_weights_after_loading() |
| 562 | logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) |
| 563 | return loaded_params |
| 564 | |
| 565 | def _finalize_moe(self) -> None: |
| 566 | for layer in self.model.layers: |
| 567 | layer.ffn.finalize_mega_moe_weights() |
| 568 | |
| 569 | def process_weights_after_loading(self) -> None: |
| 570 | self._finalize_moe() |
| 571 | |
| 572 | def _remap_dspark_name(self, name: str) -> str | None: |
| 573 | """Map a checkpoint ``mtp.{i}.*`` name to this model's parameter path. |
| 574 | |
| 575 | Returns None for non-mtp weights (owned by the target model). |
| 576 | """ |
| 577 | if name.startswith("context_wkv_proj."): |
| 578 | return f"model.{name}" |
| 579 | m = re.match(r"mtp\.(\d+)\.(.*)", name) |
| 580 | if m is None: |
| 581 | return None |
| 582 | stage = int(m.group(1)) |
| 583 | rest = m.group(2) |
| 584 | if rest.startswith("confidence_head.") and self.model.confidence_head is None: |
| 585 | return None |
| 586 | # Head-stack params live at model level (mtp.last), context combiner at |
| 587 | # model level (mtp.0); everything else is a per-layer decoder block. |
| 588 | head_prefixes = ( |
| 589 | "norm.", |
| 590 | "hc_head_fn", |
| 591 | "hc_head_base", |
| 592 | "hc_head_scale", |
| 593 | "markov_head.", |
| 594 | "confidence_head.", |
| 595 | ) |
| 596 | if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( |
| 597 | head_prefixes |
| 598 | ): |
| 599 | return f"model.{rest}" |
| 600 | return f"model.layers.{stage}.{rest}" |