| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch.nn as nn |
| 4 | |
| 5 | from vllm.config import VllmConfig, replace |
| 6 | from vllm.model_executor.model_loader import get_model |
| 7 | from vllm.model_executor.model_loader.utils import get_draft_load_config |
| 8 | from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( |
| 9 | _should_share, |
| 10 | get_target_lm_head, |
| 11 | maybe_share_target_embed, |
| 12 | ) |
| 13 | from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config |
| 14 | |
| 15 | |
| 16 | def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: |
| 17 | from vllm.compilation.backends import set_model_tag |
| 18 | from vllm.model_executor.models.qwen3_dflash import ( |
| 19 | dflash_has_any_non_causal, |
| 20 | ) |
| 21 | |
| 22 | speculative_config = vllm_config.speculative_config |
| 23 | assert speculative_config is not None |
| 24 | draft_model_config = speculative_config.draft_model_config |
| 25 | # Select an attention backend that supports the drafter's attention: mixing |
| 26 | # a non-causal layer onto a causal-only backend would fail. |
| 27 | draft_vllm_config = replace( |
| 28 | vllm_config, |
| 29 | attention_config=replace( |
| 30 | vllm_config.attention_config, |
| 31 | use_non_causal=dflash_has_any_non_causal(draft_model_config.hf_config), |
| 32 | backend=speculative_config.attention_backend, |
| 33 | ), |
| 34 | cache_config=( |
| 35 | replace( |
| 36 | vllm_config.cache_config, |
| 37 | cache_dtype=speculative_config.kv_cache_dtype, |
| 38 | ) |
| 39 | if speculative_config.kv_cache_dtype is not None |
| 40 | else vllm_config.cache_config |
| 41 | ), |
| 42 | load_config=get_pp_safe_draft_load_config(get_draft_load_config(vllm_config)), |
| 43 | ) |
| 44 | with set_model_tag("dflash_head"): |
| 45 | dflash_model = get_model( |
| 46 | vllm_config=draft_vllm_config, model_config=draft_model_config |
| 47 | ) |
| 48 | |
| 49 | target_language_model = ( |
| 50 | target_model.get_language_model() |
| 51 | if hasattr(target_model, "get_language_model") |
| 52 | else target_model |
| 53 | ) |
| 54 | # MuseGlimmerForCausalLM marks its inner MuseGlimmerModel as the language |
| 55 | # model, so get_language_model() already returns the inner module and has |
| 56 | # no .model of its own. |
| 57 | target_inner = getattr(target_language_model, "model", target_language_model) |
| 58 | draft_inner = dflash_model.model |
| 59 | |
| 60 | maybe_share_target_embed(dflash_model, draft_inner, target_inner) |
| 61 | |
| 62 | target_lm_head = get_target_lm_head(target_model, target_language_model) |
| 63 | draft_lm_head = getattr(dflash_model, "lm_head", None) |
| 64 | if ( |
| 65 | target_lm_head is not None |
| 66 | and not getattr(dflash_model.config, "has_own_lm_head", False) |
| 67 | and _should_share( |
| 68 | dflash_model, "has_own_lm_head", draft_lm_head, target_lm_head |
| 69 | ) |
| 70 | ): |
| 71 | if draft_lm_head is not None: |
| 72 | del dflash_model.lm_head |
| 73 | dflash_model.lm_head = target_lm_head |
| 74 | |
| 75 | return dflash_model |