| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch.nn as nn |
| 5 | |
| 6 | from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace |
| 7 | from vllm.logger import init_logger |
| 8 | from vllm.model_executor.model_loader.utils import get_draft_load_config |
| 9 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 10 | from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config |
| 11 | |
| 12 | logger = init_logger(__name__) |
| 13 | |
| 14 | |
| 15 | def _resolve_dspark_attention_backend( |
| 16 | draft_model_config: ModelConfig, |
| 17 | draft_backend: AttentionBackendEnum | None, |
| 18 | target_backend: AttentionBackendEnum | None, |
| 19 | ) -> AttentionBackendEnum | None: |
| 20 | if draft_backend is not None: |
| 21 | return draft_backend |
| 22 | # DeepSeek-V4(.1) draft layers share the target's KV-cache layout. Other |
| 23 | # DSpark architectures may use a different attention kind. |
| 24 | if draft_model_config.hf_config.model_type in ("deepseek_v4", "deepseek_v41"): |
| 25 | if target_backend is not None: |
| 26 | logger.info_once( |
| 27 | "Using the target model's %s attention backend for the " |
| 28 | "DeepSeek-V4 DSpark drafter.", |
| 29 | target_backend.name, |
| 30 | ) |
| 31 | return target_backend |
| 32 | return None |
| 33 | |
| 34 | |
| 35 | def _get_dspark_parallel_config( |
| 36 | parallel_config: ParallelConfig, |
| 37 | tensor_parallel_size: int, |
| 38 | ) -> ParallelConfig: |
| 39 | if parallel_config.enable_eplb: |
| 40 | logger.warning_once( |
| 41 | "EPLB is disabled for the DSpark draft model. EPLB remains enabled " |
| 42 | "for the target model." |
| 43 | ) |
| 44 | |
| 45 | return replace( |
| 46 | parallel_config, |
| 47 | pipeline_parallel_size=1, |
| 48 | tensor_parallel_size=tensor_parallel_size, |
| 49 | enable_eplb=False, |
| 50 | eplb_config=replace( |
| 51 | parallel_config.eplb_config, |
| 52 | num_redundant_experts=0, |
| 53 | ), |
| 54 | enable_elastic_ep=False, |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: |
| 59 | speculative_config = vllm_config.speculative_config |
| 60 | assert speculative_config is not None |
| 61 | draft_model_config = speculative_config.draft_model_config |
| 62 | |
| 63 | from vllm.compilation.backends import set_model_tag |
| 64 | from vllm.model_executor.model_loader import get_model |
| 65 | from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal |
| 66 | from vllm.model_executor.models.utils import get_draft_quant_config |
| 67 | from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( |
| 68 | _should_share, |
| 69 | get_target_lm_head, |
| 70 | maybe_share_target_embed, |
| 71 | ) |
| 72 | |
| 73 | draft_attention_backend = _resolve_dspark_attention_backend( |
| 74 | draft_model_config, |
| 75 | speculative_config.attention_backend, |
| 76 | vllm_config.attention_config.backend, |
| 77 | ) |
| 78 | |
| 79 | draft_vllm_config = replace( |
| 80 | vllm_config, |
| 81 | parallel_config=_get_dspark_parallel_config( |
| 82 | vllm_config.parallel_config, |
| 83 | speculative_config.draft_parallel_config.tensor_parallel_size, |
| 84 | ), |
| 85 | attention_config=replace( |
| 86 | vllm_config.attention_config, |
| 87 | use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), |
| 88 | backend=draft_attention_backend, |
| 89 | ), |
| 90 | cache_config=( |
| 91 | replace( |
| 92 | vllm_config.cache_config, |
| 93 | cache_dtype=speculative_config.kv_cache_dtype, |
| 94 | ) |
| 95 | if speculative_config.kv_cache_dtype is not None |
| 96 | else vllm_config.cache_config |
| 97 | ), |
| 98 | load_config=get_pp_safe_draft_load_config(get_draft_load_config(vllm_config)), |
| 99 | ) |
| 100 | # VllmConfig post-init restores the target's quant config because the target |
| 101 | # config is retained for DSpark's target-layer metadata, so we must override it. |
| 102 | draft_vllm_config.quant_config = get_draft_quant_config(vllm_config) |
| 103 | |
| 104 | with set_model_tag("dspark_head"): |
| 105 | draft_model = get_model( |
| 106 | vllm_config=draft_vllm_config, model_config=draft_model_config |
| 107 | ) |
| 108 | |
| 109 | target_language_model = ( |
| 110 | target_model.get_language_model() |
| 111 | if hasattr(target_model, "get_language_model") |
| 112 | else target_model |
| 113 | ) |
| 114 | target_inner = target_language_model.model |
| 115 | draft_inner = draft_model.model |
| 116 | target_vocab_size = vllm_config.model_config.get_vocab_size() |
| 117 | |
| 118 | if draft_model_config.get_vocab_size() <= target_vocab_size: |
| 119 | maybe_share_target_embed(draft_model, draft_inner, target_inner) |
| 120 | |
| 121 | target_lm_head = get_target_lm_head(target_model, target_language_model) |
| 122 | draft_lm_head = getattr(draft_model, "lm_head", None) |
| 123 | draft_output_vocab_size = ( |
| 124 | getattr(draft_model_config.hf_config, "draft_vocab_size", None) |
| 125 | or draft_model_config.get_vocab_size() |
| 126 | ) |
| 127 | if ( |
| 128 | target_lm_head is not None |
| 129 | and draft_output_vocab_size == target_vocab_size |
| 130 | and _should_share(draft_model, "has_own_lm_head", draft_lm_head, target_lm_head) |
| 131 | ): |
| 132 | if draft_lm_head is not None: |
| 133 | del draft_model.lm_head |
| 134 | draft_model.lm_head = target_lm_head |
| 135 | |
| 136 | return draft_model |