@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e439c784fa

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 13, 03:15 UTC

Add support for Eagle with separate lm-head and embed_tokens layers (#28549)

Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com>

Eldar Kurtić committed Nov 15, 2025 at 15:12 UTC e439c784fa318dbc23c04b0730bee0fccf46481d

12 files changed +204 -63

tests/v1/spec_decode/test_eagle.py

+17 -16

@@ -324,6 +324,7 @@ def test_prepare_inputs_padded():

324 @pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform())

325 @pytest.mark.parametrize("pp_size", [1, 2])

326 @pytest.mark.parametrize("use_distinct_embed_tokens", [True, False])

327 +@pytest.mark.parametrize("use_distinct_lm_head", [True, False])

328 @mock.patch("vllm.v1.spec_decode.eagle.get_pp_group")

329 @mock.patch("vllm.v1.spec_decode.eagle.get_layers_from_vllm_config")

330 @mock.patch("vllm.v1.spec_decode.eagle.get_model")

@@ -335,6 +336,7 @@ def test_load_model(

336 attn_backend,

337 pp_size,

338 use_distinct_embed_tokens,

339 + use_distinct_lm_head,

340 monkeypatch,

341 ):

342 monkeypatch.setenv("VLLM_ATTENTION_BACKEND", attn_backend)

@@ -350,12 +352,13 @@ def test_load_model(

352

353 # Setup draft model mock

354 mock_model = mock.MagicMock()

355 + mock_model.model = mock.MagicMock()

356 + mock_model.has_own_embed_tokens = use_distinct_embed_tokens

357 if use_distinct_embed_tokens:

354 - # Some models can have a different hidden size than the target model,

355 - # so we test that their embed_tokens doesn't get overwritten

356 - mock_model.model.embed_tokens.weight.shape = (131072, 2048)

357 - else:

358 - mock_model.model.embed_tokens.weight.shape = (131072, 4096)

358 + mock_model.model.embed_tokens = mock.MagicMock()

359 + mock_model.has_own_lm_head = use_distinct_lm_head

360 + if use_distinct_lm_head:

361 + mock_model.lm_head = mock.MagicMock()

362

363 mock_get_model.return_value = mock_model

364

@@ -391,15 +394,13 @@ def test_load_model(

394

395 target_model = mock.create_autospec(_TargetModelStub, instance=True)

396 target_model.model = mock.MagicMock()

394 - target_model.model.embed_tokens.weight.shape = (131072, 4096)

397 + target_model.lm_head = mock.MagicMock()

398 + target_model.model.embed_tokens = mock.MagicMock()

399

400 from vllm.model_executor.models import SupportsMultiModal

401

402 assert not isinstance(target_model, SupportsMultiModal)

403

400 - if method == "eagle":

401 - target_model.lm_head = mock.MagicMock()

402 -

404 # Create proposer using the helper function

405 proposer = _create_proposer(method, num_speculative_tokens=8)

406

@@ -409,18 +410,18 @@ def test_load_model(

410 # Verify common interactions

411 mock_get_model.assert_called_once()

412

412 - # Verify that EAGLE models gain the lm head from the target model

413 - if method == "eagle":

414 - assert proposer.model.lm_head == target_model.lm_head

413 + # Verify that the lm head is set correctly

414 + if use_distinct_lm_head:

415 + assert proposer.model.lm_head is not target_model.lm_head

416 + else:

417 + assert proposer.model.lm_head is target_model.lm_head

418

419 # Verify that the embed tokens are set correctly

420 # If pp_size is > 1, the embed tokens should be distinct

421 if pp_size > 1 or use_distinct_embed_tokens:

419 - assert proposer.model.model.embed_tokens != target_model.model.embed_tokens

422 + assert proposer.model.model.embed_tokens is not target_model.model.embed_tokens

423 else:

421 - # When pp_size is 1 and the draft and target models have

422 - # embed_tokens of the same shape, they should be shared.

423 - assert proposer.model.model.embed_tokens == target_model.model.embed_tokens

424 + assert proposer.model.model.embed_tokens is target_model.model.embed_tokens

425

426

427 @pytest.mark.parametrize("method", ["eagle", "eagle3"])

tests/v1/spec_decode/test_mtp.py

@@ -67,6 +67,10 @@ def test_mtp_load_model_unified(mock_get_model, mock_get_layers, mock_get_pp_gro

67 mock_model = mock.MagicMock()

68 mock_model.model.embed_tokens.weight.shape = (131072, 4096)

69 mock_get_model.return_value = mock_model

70 + # MTP does not have its own embed_tokens or lm_head

71 + # so it should share them with the target model

72 + mock_model.has_own_embed_tokens = False

73 + mock_model.has_own_lm_head = False

75 target_attn_layers = {"target_attn_1": mock.MagicMock()}

76 all_attn_layers = {**target_attn_layers, "draft_attn_1": mock.MagicMock()}

vllm/model_executor/models/deepseek_eagle.py

+2 -1

@@ -26,7 +26,7 @@ from vllm.model_executor.models.deepseek_v2 import (

26 )

27 from vllm.utils import init_logger

29 -from .utils import AutoWeightsLoader, maybe_prefix

29 +from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight

31 logger = init_logger(__name__)

@@ -250,6 +250,7 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):

250 name, loaded_weight = inputs

251 if "lm_head" not in name:

252 name = "model." + name

253 + process_eagle_weight(self, name)

254 return name, loaded_weight

255

256 loader = AutoWeightsLoader(

vllm/model_executor/models/deepseek_v2.py

+2 -2

@@ -85,7 +85,7 @@ from vllm.v1.attention.backends.mla.indexer import (

85 )

86 from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec

88 -from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP

88 +from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP

89 from .utils import (

90 PPMissingLayer,

91 is_pp_missing_parameter,

@@ -1311,7 +1311,7 @@ class DeepseekV2MixtureOfExperts(MixtureOfExperts):

1311

1312

1313 class DeepseekV2ForCausalLM(

1314 - nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA

1314 + nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA, SupportsEagle

1315 ):

1316 packed_modules_mapping = {

1317 "gate_up_proj": ["gate_proj", "up_proj"],

vllm/model_executor/models/interfaces.py

+65 -5

@@ -932,13 +932,73 @@ def supports_transcription(

932

933

934 @runtime_checkable

935 -class SupportsEagle3(Protocol):

935 +class SupportsEagleBase(Protocol):

936 + """Base interface for models that support EAGLE-based speculative decoding."""

937 +

938 + has_own_lm_head: bool = False

939 + """

940 + A flag that indicates this model has trained its own lm_head.

941 + """

942 +

943 + has_own_embed_tokens: bool = False

944 + """

945 + A flag that indicates this model has trained its own input embeddings.

946 + """

947 +

948 +

949 +@overload

950 +def supports_any_eagle(model: type[object]) -> TypeIs[type[SupportsEagleBase]]: ...

951 +

952 +

953 +@overload

954 +def supports_any_eagle(model: object) -> TypeIs[SupportsEagleBase]: ...

955 +

956 +

957 +def supports_any_eagle(

958 + model: type[object] | object,

959 +) -> TypeIs[type[SupportsEagleBase]] | TypeIs[SupportsEagleBase]:

960 + """Check if model supports any EAGLE variant (1, 2, or 3)."""

961 + return supports_eagle(model) or supports_eagle3(model)

962 +

963 +

964 +@runtime_checkable

965 +class SupportsEagle(SupportsEagleBase, Protocol):

966 + """The interface required for models that support

967 + EAGLE-1 and EAGLE-2 speculative decoding."""

968 +

969 + supports_eagle: ClassVar[Literal[True]] = True

970 + """

971 + A flag that indicates this model supports EAGLE-1 and EAGLE-2

972 + speculative decoding.

973 +

974 + Note:

975 + There is no need to redefine this flag if this class is in the

976 + MRO of your model class.

977 + """

978 +

979 +

980 +@overload

981 +def supports_eagle(model: type[object]) -> TypeIs[type[SupportsEagle]]: ...

982 +

983 +

984 +@overload

985 +def supports_eagle(model: object) -> TypeIs[SupportsEagle]: ...

986 +

987 +

988 +def supports_eagle(

989 + model: type[object] | object,

990 +) -> TypeIs[type[SupportsEagle]] | TypeIs[SupportsEagle]:

991 + return isinstance(model, SupportsEagle)

992 +

993 +

994 +@runtime_checkable

995 +class SupportsEagle3(SupportsEagleBase, Protocol):

996 """The interface required for models that support

937 - EAGLE3 speculative decoding."""

997 + EAGLE-3 speculative decoding."""

998

999 supports_eagle3: ClassVar[Literal[True]] = True

1000 """

941 - A flag that indicates this model supports EAGLE3

1001 + A flag that indicates this model supports EAGLE-3

1002 speculative decoding.

1003

1004 Note:

@@ -949,7 +1009,7 @@ class SupportsEagle3(Protocol):

1009 def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:

1010 """

1011 Set which layers should output auxiliary

952 - hidden states for EAGLE3.

1012 + hidden states for EAGLE-3.

1013

1014 Args:

1015 layers: Tuple of layer indices that should output auxiliary

@@ -960,7 +1020,7 @@ class SupportsEagle3(Protocol):

1020 def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:

1021 """

1022 Get the layer indices that should output auxiliary hidden states

963 - for EAGLE3.

1023 + for EAGLE-3.

1024

1025 Returns:

1026 Tuple of layer indices for auxiliary hidden state outputs.

vllm/model_executor/models/llama.py

+4 -2

@@ -58,7 +58,7 @@ from vllm.model_executor.model_loader.weight_utils import (

58 )

59 from vllm.sequence import IntermediateTensors

61 -from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP

61 +from .interfaces import SupportsEagle, SupportsEagle3, SupportsLoRA, SupportsPP

62 from .utils import (

63 AutoWeightsLoader,

64 PPMissingLayer,

@@ -529,7 +529,9 @@ class LlamaModel(nn.Module):

529 return loaded_params

530

531

532 -class LlamaForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

532 +class LlamaForCausalLM(

533 + nn.Module, SupportsLoRA, SupportsPP, SupportsEagle, SupportsEagle3

534 +):

535 packed_modules_mapping = {

536 "qkv_proj": ["q_proj", "k_proj", "v_proj"],

537 "gate_up_proj": ["gate_proj", "up_proj"],

vllm/model_executor/models/llama4_eagle.py

+2 -1

@@ -35,7 +35,7 @@ from vllm.model_executor.models.llama4 import Llama4DecoderLayer, Llama4ForCausa

35 from vllm.model_executor.models.utils import extract_layer_index

37 from .interfaces import SupportsMultiModal

38 -from .utils import AutoWeightsLoader, maybe_prefix

38 +from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight

40 logger = init_logger(__name__)

@@ -212,6 +212,7 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):

212 name, weight = self.permute_qk_weight_for_rotary(name, loaded_weight)

213 if "lm_head" not in name:

214 name = "model." + name

215 + process_eagle_weight(self, name)

216 return name, weight

217

218 loader = AutoWeightsLoader(

vllm/model_executor/models/llama_eagle.py

+2 -1

@@ -17,7 +17,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmb

17 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

18 from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM

20 -from .utils import AutoWeightsLoader, maybe_prefix

20 +from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight

22 logger = init_logger(__name__)

@@ -179,6 +179,7 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):

179 name, loaded_weight = inputs

180 if "lm_head" not in name:

181 name = "model." + name

182 + process_eagle_weight(self, name)

183 return name, loaded_weight

184

185 loader = AutoWeightsLoader(

vllm/model_executor/models/llama_eagle3.py

+2 -1

@@ -23,7 +23,7 @@ from vllm.model_executor.models.llama import LlamaDecoderLayer, LlamaForCausalLM

23 from vllm.multimodal import MULTIMODAL_REGISTRY

24 from vllm.multimodal.inputs import NestedTensors

26 -from .utils import AutoWeightsLoader, maybe_prefix

26 +from .utils import AutoWeightsLoader, maybe_prefix, process_eagle_weight

28 logger = init_logger(__name__)

@@ -324,6 +324,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

324 if "embed_tokens" in name:

325 includes_embed_tokens = True

326 model_weights[name] = loaded_weight

327 + process_eagle_weight(self, name)

328

329 skip_substrs = []

330 if not includes_draft_id_mapping:

vllm/model_executor/models/minicpm_eagle.py

+9 -3

@@ -43,7 +43,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (

43 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

44 from vllm.sequence import IntermediateTensors

46 -from .interfaces import SupportsLoRA, SupportsPP

46 +from .interfaces import SupportsEagle, SupportsLoRA, SupportsPP

47 from .minicpm import MiniCPMAttention as EagleMiniCPMAttention

48 from .minicpm import MiniCPMMLP as EagleMiniCPMMLP

49 from .minicpm import MiniCPMMoE as EagleMiniCPMMoE

@@ -52,6 +52,7 @@ from .utils import (

52 is_pp_missing_parameter,

53 make_empty_intermediate_tensors_factory,

54 maybe_prefix,

55 + process_eagle_weight,

56 )

@@ -289,7 +290,7 @@ class EagleMiniCPMModel(nn.Module):

290 return loaded_params

291

292

292 -class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

293 +class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle):

294 packed_modules_mapping = {

295 "qkv_proj": [

296 "q_proj",

@@ -376,8 +377,13 @@ class EagleMiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

377 return logits

378

379 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

380 + def transform(inputs):

381 + name, loaded_weight = inputs

382 + process_eagle_weight(self, name)

383 + return name, loaded_weight

384 +

385 loader = AutoWeightsLoader(

386 self,

387 skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None),

388 )

383 - return loader.load_weights(weights)

389 + return loader.load_weights(map(transform, weights))

vllm/model_executor/models/utils.py

+23

@@ -19,6 +19,7 @@ from vllm.distributed import (

19 )

20 from vllm.logger import init_logger

21 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

22 +from vllm.model_executor.models.interfaces import supports_any_eagle

23 from vllm.multimodal import NestedTensors

24 from vllm.sequence import IntermediateTensors

25 from vllm.utils.math_utils import cdiv

@@ -825,3 +826,25 @@ direct_register_custom_op(

826 fake_impl=sequence_parallel_chunk_impl_fake,

827 tags=(torch.Tag.needs_fixed_stride_order,),

828 )

829 +

830 +

831 +def process_eagle_weight(

832 + model: nn.Module,

833 + name: str,

834 +) -> None:

835 + """

836 + Update EAGLE model flags based on loaded weight name.

837 + This should be called during weight loading to detect if a model

838 + has its own lm_head or embed_tokens weight.

839 + Args:

840 + model: The model instance (must support EAGLE)

841 + name: The name of the weight to process

842 + """

843 + if not supports_any_eagle(model):

844 + return

845 +

846 + # To prevent overriding with target model's layers

847 + if "lm_head" in name:

848 + model.has_own_lm_head = True

849 + if "embed_tokens" in name:

850 + model.has_own_embed_tokens = True

vllm/v1/spec_decode/eagle.py

+72 -31

@@ -991,6 +991,7 @@ class EagleProposer:

991 target_language_model = target_model.get_language_model()

992 else:

993 target_language_model = target_model

994 +

995 # share embed_tokens with the target model if needed

996 if get_pp_group().world_size == 1:

997 if hasattr(target_language_model.model, "embed_tokens"):

@@ -1002,52 +1003,92 @@ class EagleProposer:

1003 "Target model does not have 'embed_tokens' or 'embedding' attribute"

1004 )

1005

1005 - # Check if shapes match and we found the embedding

1006 - eagle_shape = self.model.model.embed_tokens.weight.shape

1007 - target_shape = target_embed_tokens.weight.shape

1008 - if eagle_shape == target_shape:

1009 - logger.info(

1010 - "Assuming the EAGLE head shares the same vocab embedding"

1011 - " with the target model."

1012 - )

1013 - del self.model.model.embed_tokens

1014 - self.model.model.embed_tokens = target_embed_tokens

1006 + share_embeddings = False

1007 + if hasattr(self.model, "has_own_embed_tokens"):

1008 + # EAGLE model

1009 + if not self.model.has_own_embed_tokens:

1010 + share_embeddings = True

1011 + logger.info(

1012 + "Detected EAGLE model without its own embed_tokens in the"

1013 + " checkpoint. Sharing target model embedding weights with the"

1014 + " draft model."

1015 + )

1016 + elif (

1017 + isinstance(target_embed_tokens.weight, torch.Tensor)

1018 + and isinstance(self.model.model.embed_tokens.weight, torch.Tensor)

1019 + and torch.equal(

1020 + target_embed_tokens.weight, self.model.model.embed_tokens.weight

1021 + )

1022 + ):

1023 + share_embeddings = True

1024 + logger.info(

1025 + "Detected EAGLE model with embed_tokens identical to the target"

1026 + " model. Sharing target model embedding weights with the draft"

1027 + " model."

1028 + )

1029 + else:

1030 + logger.info(

1031 + "Detected EAGLE model with distinct embed_tokens weights. "

1032 + "Keeping separate embedding weights from the target model."

1033 + )

1034 else:

1035 + # MTP model

1036 + share_embeddings = True

1037 logger.info(

1017 - "The EAGLE head's vocab embedding will be loaded separately"

1018 - " from the target model."

1038 + "Detected MTP model. "

1039 + "Sharing target model embedding weights with the draft model."

1040 )

1041 +

1042 + if share_embeddings:

1043 + if hasattr(self.model.model, "embed_tokens"):

1044 + del self.model.model.embed_tokens

1045 + self.model.model.embed_tokens = target_embed_tokens

1046 else:

1047 logger.info(

1022 - "The EAGLE head's vocab embedding will be loaded separately"

1048 + "The draft model's vocab embedding will be loaded separately"

1049 " from the target model."

1050 )

1051

1052 # share lm_head with the target model if needed

1027 - # some model definition do not define lm_head explicitly

1028 - # and reuse embed_tokens for lm_head, e.g., CohereForCausalLM

1029 - if self.vllm_config.speculative_config.method != "eagle3":

1030 - if hasattr(target_language_model, "lm_head"):

1031 - logger.info("Loading EAGLE LM head weights from the target model.")

1032 - self.model.lm_head = target_language_model.lm_head

1033 - else:

1034 - if (

1035 - hasattr(self.model, "lm_head")

1036 - and hasattr(target_language_model, "lm_head")

1037 - and self.model.lm_head.weight.shape

1038 - == target_language_model.lm_head.weight.shape

1053 + share_lm_head = False

1054 + if hasattr(self.model, "has_own_lm_head"):

1055 + # EAGLE model

1056 + if not self.model.has_own_lm_head:

1057 + share_lm_head = True

1058 + logger.info(

1059 + "Detected EAGLE model without its own lm_head in the checkpoint. "

1060 + "Sharing target model lm_head weights with the draft model."

1061 + )

1062 + elif (

1063 + hasattr(target_language_model, "lm_head")

1064 + and isinstance(target_language_model.lm_head.weight, torch.Tensor)

1065 + and isinstance(self.model.lm_head.weight, torch.Tensor)

1066 + and torch.equal(

1067 + target_language_model.lm_head.weight, self.model.lm_head.weight

1068 + )

1069 ):

1070 + share_lm_head = True

1071 logger.info(

1041 - "Assuming the EAGLE head shares the same lm_head"

1042 - " with the target model."

1072 + "Detected EAGLE model with lm_head identical to the target model. "

1073 + "Sharing target model lm_head weights with the draft model."

1074 )

1044 - del self.model.lm_head

1045 - self.model.lm_head = target_language_model.lm_head

1075 else:

1076 logger.info(

1048 - "The EAGLE head's lm_head will be loaded separately"

1049 - " from the target model."

1077 + "Detected EAGLE model with distinct lm_head weights. "

1078 + "Keeping separate lm_head weights from the target model."

1079 )

1080 + else:

1081 + # MTP model

1082 + share_lm_head = True

1083 + logger.info(

1084 + "Detected MTP model. "

1085 + "Sharing target model lm_head weights with the draft model."

1086 + )

1087 +

1088 + if share_lm_head and hasattr(target_language_model, "lm_head"):

1089 + if hasattr(self.model, "lm_head"):

1090 + del self.model.lm_head

1091 + self.model.lm_head = target_language_model.lm_head

1092

1093 @torch.inference_mode()

1094 def dummy_run(