| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from transformers import DeepseekV3Config |
| 5 | |
| 6 | |
| 7 | class K3DSparkConfig(DeepseekV3Config): |
| 8 | """Configuration for a dense MLA DSpark draft model.""" |
| 9 | |
| 10 | model_type = "k3_dspark" |
| 11 | has_no_defaults_at_init = True |
| 12 | |
| 13 | def __init__( |
| 14 | self, |
| 15 | mla_use_nope: bool = False, |
| 16 | mla_use_output_gate: bool = False, |
| 17 | mla_use_qk_norm: bool = False, |
| 18 | rope_theta: float = 50000.0, |
| 19 | **kwargs, |
| 20 | ) -> None: |
| 21 | # DeepseekV3Config defaults to a MoE topology. Zero these fields so |
| 22 | # generic vLLM config logic also recognizes this draft as dense. |
| 23 | kwargs.setdefault("n_routed_experts", 0) |
| 24 | kwargs.setdefault("n_shared_experts", 0) |
| 25 | kwargs.setdefault("num_experts_per_tok", 0) |
| 26 | |
| 27 | rope_parameters = kwargs.get("rope_parameters") |
| 28 | if rope_parameters is None: |
| 29 | kwargs["rope_parameters"] = { |
| 30 | "rope_type": "default", |
| 31 | "rope_theta": rope_theta, |
| 32 | } |
| 33 | else: |
| 34 | rope_parameters = dict(rope_parameters) |
| 35 | rope_parameters.setdefault("rope_type", "default") |
| 36 | rope_parameters.setdefault("rope_theta", rope_theta) |
| 37 | kwargs["rope_parameters"] = rope_parameters |
| 38 | |
| 39 | super().__init__(**kwargs) |
| 40 | self.mla_use_nope = mla_use_nope |
| 41 | self.mla_use_output_gate = mla_use_output_gate |
| 42 | self.mla_use_qk_norm = mla_use_qk_norm |
| 43 | |
| 44 | unsupported = [ |
| 45 | name |
| 46 | for name in ( |
| 47 | "mla_use_nope", |
| 48 | "mla_use_output_gate", |
| 49 | "mla_use_qk_norm", |
| 50 | "dspark_bonus_anchor", |
| 51 | ) |
| 52 | if getattr(self, name, False) |
| 53 | ] |
| 54 | if self.q_lora_rank is None: |
| 55 | unsupported.append("q_lora_rank=None") |
| 56 | if unsupported: |
| 57 | raise ValueError( |
| 58 | "MLA DSpark does not support " + ", ".join(unsupported) + "." |
| 59 | ) |
| 60 | |
| 61 | self.draft_vocab_size = ( |
| 62 | getattr(self, "draft_vocab_size", None) or self.vocab_size |
| 63 | ) |
| 64 | if self.draft_vocab_size != self.vocab_size: |
| 65 | raise ValueError( |
| 66 | "MLA DSpark requires draft_vocab_size to equal vocab_size when " |
| 67 | "sharing the target embedding and LM head." |
| 68 | ) |
| 69 | |
| 70 | target_layer_ids = getattr(self, "target_layer_ids", None) |
| 71 | if not target_layer_ids or getattr(self, "num_target_layers", None) != len( |
| 72 | target_layer_ids |
| 73 | ): |
| 74 | raise ValueError( |
| 75 | "MLA DSpark requires non-empty target_layer_ids and a matching " |
| 76 | "num_target_layers." |
| 77 | ) |