8b346309a5 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py

+2 -2

@@ -13,15 +13,15 @@ import torch

13 import torch.nn as nn

14

15 from vllm.config import VllmConfig

16 +from vllm.model_executor.models.interfaces import EagleModelMixin

17 from vllm.model_executor.models.llama import LlamaForCausalLM

18 from vllm.sequence import IntermediateTensors

19

20

20 -class PredictableLlamaModel(nn.Module):

21 +class PredictableLlamaModel(nn.Module, EagleModelMixin):

22 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

23 super().__init__()

24 self.config = vllm_config.model_config.hf_config

24 - self.aux_hidden_state_layers = tuple[int, ...]()

25

26 # Create minimal embed_tokens for embedding

27 from vllm.model_executor.layers.vocab_parallel_embedding import (

vllm/model_executor/models/afmoe.py

+6 -15

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

37 maybe_remap_kv_scale_name,

38 )

39 from vllm.model_executor.models.interfaces import (

40 + EagleModelMixin,

41 SupportsEagle3,

42 SupportsLoRA,

43 SupportsPP,

@@ -384,7 +385,7 @@ class AfmoeDecoderLayer(nn.Module):

385 "inputs_embeds": 0,

386 }

387 )

387 -class AfmoeModel(nn.Module):

388 +class AfmoeModel(nn.Module, EagleModelMixin):

389 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

390 super().__init__()

391

@@ -421,8 +422,6 @@ class AfmoeModel(nn.Module):

422 else:

423 self.norm = PPMissingLayer()

424

424 - self.aux_hidden_state_layers = tuple[int, ...]()

425 -

425 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

426 ["hidden_states", "residual"], config.hidden_size

427 )

@@ -453,15 +452,14 @@ class AfmoeModel(nn.Module):

452 hidden_states = intermediate_tensors["hidden_states"]

453 residual = intermediate_tensors["residual"]

454

456 - aux_hidden_states = []

455 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

456 for idx, layer in enumerate(

457 islice(self.layers, self.start_layer, self.end_layer)

458 ):

460 - if idx in self.aux_hidden_state_layers:

461 - aux_hidden_states.append(

462 - hidden_states + residual if residual is not None else hidden_states

463 - )

459 hidden_states, residual = layer(positions, hidden_states, residual)

460 + self._maybe_add_hidden_state(

461 + aux_hidden_states, idx + 1, hidden_states, residual

462 + )

463

464 if not get_pp_group().is_last_rank:

465 return IntermediateTensors(

@@ -691,13 +689,6 @@ class AfmoeForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

689 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

690 return self.model.embed_input_ids(input_ids)

691

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

695 - self.model.aux_hidden_state_layers = layers

696 -

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

698 - num_layers = len(self.model.layers)

699 - return (2, num_layers // 2, num_layers - 3)

700 -

692 def forward(

693 self,

694 input_ids: torch.Tensor | None,

vllm/model_executor/models/apertus.py

+15 -15

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

60 from vllm.sequence import IntermediateTensors

61 from vllm.v1.attention.backend import AttentionType

62

63 -from .interfaces import SupportsLoRA, SupportsPP

63 +from .interfaces import (

64 + EagleModelMixin,

65 + SupportsEagle,

66 + SupportsEagle3,

67 + SupportsLoRA,

68 + SupportsPP,

69 +)

70 from .utils import (

71 AutoWeightsLoader,

72 PPMissingLayer,

@@ -313,7 +319,7 @@ class ApertusDecoderLayer(nn.Module):

319

320

321 @support_torch_compile

316 -class ApertusModel(nn.Module):

322 +class ApertusModel(nn.Module, EagleModelMixin):

323 def __init__(

324 self,

325 *,

@@ -357,8 +363,6 @@ class ApertusModel(nn.Module):

363 else:

364 self.norm = PPMissingLayer()

365

360 - self.aux_hidden_state_layers = tuple[int, ...]()

361 -

366 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

367 ["hidden_states", "residual"], config.hidden_size

368 )

@@ -384,13 +388,14 @@ class ApertusModel(nn.Module):

388 hidden_states = intermediate_tensors["hidden_states"]

389 residual = intermediate_tensors["residual"]

390

387 - aux_hidden_states = []

391 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

392 for idx, layer in enumerate(

393 islice(self.layers, self.start_layer, self.end_layer)

394 ):

391 - if idx in self.aux_hidden_state_layers:

392 - aux_hidden_states.append(hidden_states + residual)

395 hidden_states, residual = layer(positions, hidden_states, residual)

396 + self._maybe_add_hidden_state(

397 + aux_hidden_states, idx + 1, hidden_states, residual

398 + )

399

400 if not get_pp_group().is_last_rank:

401 return IntermediateTensors(

@@ -472,7 +477,9 @@ class ApertusModel(nn.Module):

477 return loaded_params

478

479

475 -class ApertusForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

480 +class ApertusForCausalLM(

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

482 +):

483 packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}

484

485 # LoRA specific attributes

@@ -520,13 +527,6 @@ class ApertusForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

527 self.model.make_empty_intermediate_tensors

528 )

529

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

524 - self.model.aux_hidden_state_layers = layers

525 -

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

527 - num_layers = len(self.model.layers)

528 - return (2, num_layers // 2, num_layers - 3)

529 -

530 def _init_model(

531 self,

532 vllm_config: VllmConfig,

vllm/model_executor/models/arcee.py

+15 -12

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

32 )

33 from vllm.sequence import IntermediateTensors

34

35 -from .interfaces import SupportsLoRA, SupportsPP

35 +from .interfaces import (

36 + EagleModelMixin,

37 + SupportsEagle,

38 + SupportsEagle3,

39 + SupportsLoRA,

40 + SupportsPP,

41 +)

42 from .utils import (

43 AutoWeightsLoader,

44 PPMissingLayer,

@@ -170,7 +176,7 @@ class ArceeDecoderLayer(nn.Module):

176

177

178 @support_torch_compile

173 -class ArceeModel(nn.Module):

179 +class ArceeModel(nn.Module, EagleModelMixin):

180 """The transformer model backbone for Arcee (embedding layer + stacked

181 decoder blocks + final norm)."""

182

@@ -218,10 +224,6 @@ class ArceeModel(nn.Module):

224 else:

225 self.norm = PPMissingLayer()

226

221 - # For optional capturing of intermediate hidden states

222 - # (not used by default)

223 - self.aux_hidden_state_layers: tuple[int, ...] = tuple()

224 -

227 # Prepare factory for empty intermediate tensors

228 # (for pipeline scheduling)

229 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

@@ -253,15 +255,14 @@ class ArceeModel(nn.Module):

255 hidden_states = intermediate_tensors["hidden_states"]

256 residual = intermediate_tensors["residual"]

257

256 - aux_hidden_states: list[torch.Tensor] = []

258 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

259 for idx, layer in enumerate(

260 islice(self.layers, self.start_layer, self.end_layer)

261 ):

260 - if idx in self.aux_hidden_state_layers:

261 - aux_hidden_states.append(

262 - hidden_states + residual

263 - ) # capture pre-layer hidden state if needed

262 hidden_states, residual = layer(positions, hidden_states, residual)

263 + self._maybe_add_hidden_state(

264 + aux_hidden_states, idx + 1, hidden_states, residual

265 + )

266

267 if not get_pp_group().is_last_rank:

268 # Send intermediate results to the next pipeline stage

@@ -348,7 +349,9 @@ class ArceeModel(nn.Module):

349 return loaded_params

350

351

351 -class ArceeForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

352 +class ArceeForCausalLM(

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

354 +):

355 """Arcee Model for causal language modeling, integrated with vLLM

356 runtime."""

357

vllm/model_executor/models/gpt_oss.py

+15 -14

@@ -47,7 +47,13 @@ from vllm.sequence import IntermediateTensors

47 from vllm.utils.math_utils import cdiv

48 from vllm.v1.attention.backend import AttentionType

49

50 -from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP

50 +from .interfaces import (

51 + EagleModelMixin,

52 + SupportsEagle,

53 + SupportsEagle3,

54 + SupportsLoRA,

55 + SupportsPP,

56 +)

57 from .utils import (

58 AutoWeightsLoader,

59 WeightsMapper,

@@ -256,7 +262,7 @@ class TransformerBlock(torch.nn.Module):

262

263

264 @support_torch_compile

259 -class GptOssModel(nn.Module):

265 +class GptOssModel(nn.Module, EagleModelMixin):

266 def __init__(

267 self,

268 *,

@@ -285,7 +291,6 @@ class GptOssModel(nn.Module):

291 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

292 ["hidden_states", "residual"], self.config.hidden_size

293 )

288 - self.aux_hidden_state_layers = tuple[int, ...]()

294

295 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

296 return self.embedding(input_ids)

@@ -309,12 +314,13 @@ class GptOssModel(nn.Module):

314 x = intermediate_tensors["hidden_states"]

315 residual = intermediate_tensors["residual"]

316

312 - aux_hidden_states = []

317 + aux_hidden_states = self._maybe_add_hidden_state(

318 + [], self.start_layer, x, residual

319 + )

320 for i in range(self.start_layer, self.end_layer):

321 layer = self.layers[i]

315 - if i in self.aux_hidden_state_layers:

316 - aux_hidden_states.append(x if residual is None else x + residual)

322 x, residual = layer(x, positions, residual)

323 + self._maybe_add_hidden_state(aux_hidden_states, i + 1, x, residual)

324 if not get_pp_group().is_last_rank:

325 return IntermediateTensors({"hidden_states": x, "residual": residual})

326 x, _ = self.norm(x, residual)

@@ -1141,7 +1147,9 @@ class GptOssModel(nn.Module):

1147 )

1148

1149

1144 -class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

1150 +class GptOssForCausalLM(

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

1152 +):

1153 is_3d_moe_weight: bool = True

1154 packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}

1155

@@ -1197,13 +1205,6 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

1205 self.model.make_empty_intermediate_tensors

1206 )

1207

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

1201 - self.model.aux_hidden_state_layers = layers

1202 -

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

1204 - num_layers = len(self.model.layers)

1205 - return (2, num_layers // 2, num_layers - 3)

1206 -

1208 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

1209 return self.model.embed_input_ids(input_ids)

1210

vllm/model_executor/models/hunyuan_v1.py

+17 -15

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

66 from vllm.sequence import IntermediateTensors

67 from vllm.v1.attention.backend import AttentionType

68

69 -from .interfaces import MixtureOfExperts, SupportsEagle3, SupportsLoRA, SupportsPP

69 +from .interfaces import (

70 + EagleModelMixin,

71 + MixtureOfExperts,

72 + SupportsEagle,

73 + SupportsEagle3,

74 + SupportsLoRA,

75 + SupportsPP,

76 +)

77 from .utils import (

78 AutoWeightsLoader,

79 PPMissingLayer,

@@ -586,7 +593,7 @@ class HunYuanDecoderLayer(nn.Module):

593 "inputs_embeds": 0,

594 }

595 )

589 -class HunYuanModel(nn.Module):

596 +class HunYuanModel(nn.Module, EagleModelMixin):

597 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

598 super().__init__()

599

@@ -629,7 +636,6 @@ class HunYuanModel(nn.Module):

636 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

637 else:

638 self.norm = PPMissingLayer()

632 - self.aux_hidden_state_layers = tuple[int, ...]()

639

640 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

641 return self.embed_tokens(input_ids)

@@ -654,13 +660,10 @@ class HunYuanModel(nn.Module):

660

661 cla_factor = _get_cla_factor(self.config)

662 prev_kv_states = None

657 - aux_hidden_states = []

663 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

664 for i, layer in enumerate(

665 islice(self.layers, self.start_layer, self.end_layer)

666 ):

661 - if i in self.aux_hidden_state_layers:

662 - aux_hidden_states.append(hidden_states + residual)

663 -

667 hidden_states, residual, kv_states = layer(

668 positions,

669 hidden_states,

@@ -673,6 +676,10 @@ class HunYuanModel(nn.Module):

676 else:

677 prev_kv_states = None

678

679 + self._maybe_add_hidden_state(

680 + aux_hidden_states, i + 1, hidden_states, residual

681 + )

682 +

683 if not get_pp_group().is_last_rank:

684 return IntermediateTensors(

685 {"hidden_states": hidden_states, "residual": residual}

@@ -904,7 +911,9 @@ class HunYuanModel(nn.Module):

911 return loaded_params

912

913

907 -class HunyuanV1ModelBase(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

914 +class HunyuanV1ModelBase(

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

916 +):

917 packed_modules_mapping = {

918 "qkv_proj": [

919 "q_proj",

@@ -943,13 +952,6 @@ class HunyuanV1ModelBase(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

952 else:

953 self.lm_head = PPMissingLayer()

954

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

947 - self.model.aux_hidden_state_layers = layers

948 -

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

950 - num_layers = len(self.model.layers)

951 - return (2, num_layers // 2, num_layers - 3)

952 -

955 def forward(

956 self,

957 input_ids: torch.Tensor | None,

vllm/model_executor/models/hunyuan_vision.py

+2 -7

@@ -86,6 +86,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

86

87 from .interfaces import (

88 MultiModalEmbeddings,

89 + SupportsEagle,

90 SupportsEagle3,

91 SupportsLoRA,

92 SupportsMultiModal,

@@ -801,6 +802,7 @@ class HunYuanVLForConditionalGeneration(

802 SupportsPP,

803 SupportsQuant,

804 SupportsXDRoPE,

805 + SupportsEagle,

806 SupportsEagle3,

807 ):

808 # To ensure correct weight loading and mapping.

@@ -988,13 +990,6 @@ class HunYuanVLForConditionalGeneration(

990 multimodal_embeddings += tuple(image_embeddings)

991 return multimodal_embeddings

992

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

992 - self.language_model.model.aux_hidden_state_layers = layers

993 -

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

995 - num_layers = len(self.language_model.model.layers)

996 - return (2, num_layers // 2, num_layers - 3)

997 -

993 def forward(

994 self,

995 input_ids: torch.Tensor | None,

vllm/model_executor/models/interfaces.py

+50 -7

@@ -1273,6 +1273,25 @@ def supports_any_eagle(

1273 return supports_eagle(model) or supports_eagle3(model)

1274

1275

1276 +class EagleModelMixin:

1277 + aux_hidden_state_layers: tuple[int, ...] = ()

1278 +

1279 + def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:

1280 + self.aux_hidden_state_layers = layers

1281 +

1282 + def _maybe_add_hidden_state(

1283 + self,

1284 + aux_hidden_states: list[torch.Tensor],

1285 + layer_idx: int,

1286 + hidden_states: torch.Tensor,

1287 + residual: torch.Tensor,

1288 + ) -> list[torch.Tensor]:

1289 + if layer_idx in self.aux_hidden_state_layers:

1290 + value = hidden_states + residual if residual is not None else hidden_states

1291 + aux_hidden_states.append(value)

1292 + return aux_hidden_states

1293 +

1294 +

1295 @runtime_checkable

1296 class SupportsEagle(SupportsEagleBase, Protocol):

1297 """The interface required for models that support

@@ -1320,24 +1339,48 @@ class SupportsEagle3(SupportsEagleBase, Protocol):

1339

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

1341 """

1323 - Set which layers should output auxiliary

1324 - hidden states for EAGLE-3.

1342 + Set which layers should output auxiliary hidden states for EAGLE-3.

1343

1344 Args:

1345 layers: Tuple of layer indices that should output auxiliary

1346 hidden states.

1347 """

1330 - ...

1348 + parent_ref = self

1349 + if hasattr(self, "get_language_model"):

1350 + parent_ref = self.get_language_model()

1351 + elif hasattr(self, "language_model"):

1352 + parent_ref = self.language_model

1353 + assert hasattr(parent_ref, "model"), (

1354 + "Model instance must have 'model' attribute to set number of layers"

1355 + )

1356 + assert isinstance(parent_ref.model, EagleModelMixin), (

1357 + "Model instance must inherit from EagleModelMixin to set auxiliary layers"

1358 + )

1359 + parent_ref.model._set_aux_hidden_state_layers(layers)

1360

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

1361 + def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:

1362 """

1334 - Get the layer indices that should output auxiliary hidden states

1335 - for EAGLE-3.

1363 + Get the default layer indices that should output auxiliary hidden states

1364 + for EAGLE-3 for this model. Models can override this method to provide

1365 + different default layers based on their architecture, but it is encouraged

1366 + to instead include the layer specification in the model's config if possible.

1367

1368 Returns:

1369 Tuple of layer indices for auxiliary hidden state outputs.

1370 """

1340 - ...

1371 + parent_ref = self

1372 + if hasattr(self, "get_language_model"):

1373 + parent_ref = self.get_language_model()

1374 + elif hasattr(self, "language_model"):

1375 + parent_ref = self.language_model

1376 + assert hasattr(parent_ref, "model"), (

1377 + "Model instance must have 'model' attribute to get number of layers"

1378 + )

1379 + assert hasattr(parent_ref.model, "layers"), (

1380 + "Model instance must have 'layers' attribute to get number of layers"

1381 + )

1382 + num_layers = len(parent_ref.model.layers)

1383 + return (2, num_layers // 2, num_layers - 3)

1384

1385

1386 @overload

vllm/model_executor/models/llama.py

+6 -18

@@ -61,6 +61,7 @@ from vllm.v1.attention.backend import AttentionType

61

62 from .adapters import as_embedding_model, as_seq_cls_model

63 from .interfaces import (

64 + EagleModelMixin,

65 SupportsEagle,

66 SupportsEagle3,

67 SupportsLoRA,

@@ -351,7 +352,7 @@ def llama_model_invariants(

352 # mark_unbacked_dims={"input_ids": 0},

353 shape_invariants=llama_model_invariants

354 )

354 -class LlamaModel(nn.Module):

355 +class LlamaModel(nn.Module, EagleModelMixin):

356 def __init__(

357 self,

358 *,

@@ -389,8 +390,6 @@ class LlamaModel(nn.Module):

390 else:

391 self.norm = PPMissingLayer()

392

392 - self.aux_hidden_state_layers = tuple[int, ...]()

393 -

393 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

394 ["hidden_states", "residual"], config.hidden_size

395 )

@@ -417,15 +416,16 @@ class LlamaModel(nn.Module):

416 hidden_states = intermediate_tensors["hidden_states"]

417 residual = intermediate_tensors["residual"]

418

420 - aux_hidden_states = []

419 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

420 for idx, layer in enumerate(

421 islice(self.layers, self.start_layer, self.end_layer)

422 ):

424 - if idx in self.aux_hidden_state_layers:

425 - aux_hidden_states.append(hidden_states + residual)

423 hidden_states, residual = layer(

424 positions, hidden_states, residual, **extra_layer_kwargs

425 )

426 + self._maybe_add_hidden_state(

427 + aux_hidden_states, idx + 1, hidden_states, residual

428 + )

429

430 if not get_pp_group().is_last_rank:

431 return IntermediateTensors(

@@ -556,18 +556,6 @@ class LlamaForCausalLM(

556 self.model.make_empty_intermediate_tensors

557 )

558

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

560 - self.model.aux_hidden_state_layers = layers

561 -

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

563 - """Override to return default layers for Llama

564 -

565 - Note: The GPU model runner will override this with layers from

566 - the speculative config if available, providing dynamic configuration.

567 - """

568 - num_layers = len(self.model.layers)

569 - return (2, num_layers // 2, num_layers - 3)

570 -

559 def _init_model(

560 self,

561 vllm_config: VllmConfig,

vllm/model_executor/models/llava.py

+7 -8

@@ -55,6 +55,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

55 from .clip import CLIPVisionModel

56 from .interfaces import (

57 MultiModalEmbeddings,

58 + SupportsEagle,

59 SupportsEagle3,

60 SupportsLoRA,

61 SupportsMultiModal,

@@ -503,7 +504,12 @@ def init_vision_tower_for_llava(

504 dummy_inputs=LlavaDummyInputsBuilder,

505 )

506 class LlavaForConditionalGeneration(

506 - nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3

507 + nn.Module,

508 + SupportsLoRA,

509 + SupportsMultiModal,

510 + SupportsPP,

511 + SupportsEagle,

512 + SupportsEagle3,

513 ):

514 packed_modules_mapping = {

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

@@ -527,13 +533,6 @@ class LlavaForConditionalGeneration(

533

534 raise ValueError("Only image modality is supported")

535

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

531 - self.get_language_model().model.aux_hidden_state_layers = layers

532 -

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

534 - num_layers = len(self.get_language_model().model.layers)

535 - return (2, num_layers // 2, num_layers - 3)

536 -

536 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

537 super().__init__()

538

vllm/model_executor/models/mimo_v2_flash.py

-7

@@ -682,13 +682,6 @@ class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

682 self.model.make_empty_intermediate_tensors

683 )

684

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

686 - self.model.aux_hidden_state_layers = layers

687 -

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

689 - num_layers = len(self.model.layers)

690 - return (2, num_layers // 2, num_layers - 3)

691 -

685 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

686 return self.model.embed_input_ids(input_ids)

687

vllm/model_executor/models/minicpm.py

+15 -17

@@ -63,7 +63,13 @@ from vllm.model_executor.utils import set_weight_attrs

63 from vllm.platforms import current_platform

64 from vllm.sequence import IntermediateTensors

65

66 -from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP

66 +from .interfaces import (

67 + EagleModelMixin,

68 + SupportsEagle,

69 + SupportsEagle3,

70 + SupportsLoRA,

71 + SupportsPP,

72 +)

73 from .utils import (

74 AutoWeightsLoader,

75 is_pp_missing_parameter,

@@ -391,7 +397,7 @@ class MiniCPMDecoderLayer(nn.Module):

397

398

399 @support_torch_compile

394 -class MiniCPMModel(nn.Module):

400 +class MiniCPMModel(nn.Module, EagleModelMixin):

401 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

402 super().__init__()

403

@@ -413,8 +419,6 @@ class MiniCPMModel(nn.Module):

419 self._init_layers(prefix, config, cache_config, quant_config)

420 self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

421

416 - self.aux_hidden_state_layers = tuple[int, ...]()

417 -

422 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

423 ["hidden_states", "residual"], self.config.hidden_size

424 )

@@ -455,19 +459,18 @@ class MiniCPMModel(nn.Module):

459 hidden_states = intermediate_tensors["hidden_states"]

460 residual = intermediate_tensors["residual"]

461

458 - aux_hidden_states = []

462 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

463 for idx, layer in enumerate(

464 islice(self.layers, self.start_layer, self.end_layer)

465 ):

462 - if idx in self.aux_hidden_state_layers:

463 - aux_hidden_states.append(

464 - hidden_states + residual if residual is not None else hidden_states

465 - )

466 hidden_states, residual = layer(

467 positions,

468 hidden_states,

469 residual,

470 )

471 + self._maybe_add_hidden_state(

472 + aux_hidden_states, idx + 1, hidden_states, residual

473 + )

474

475 if not get_pp_group().is_last_rank:

476 return IntermediateTensors(

@@ -550,7 +553,9 @@ class MiniCPMModel(nn.Module):

553 return loaded_params

554

555

553 -class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

556 +class MiniCPMForCausalLM(

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

558 +):

559 packed_modules_mapping = {

560 "qkv_proj": [

561 "q_proj",

@@ -611,13 +616,6 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

616 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

617 return self.model.embed_input_ids(input_ids)

618

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

615 - self.model.aux_hidden_state_layers = layers

616 -

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

618 - num_layers = len(self.model.layers)

619 - return (2, num_layers // 2, num_layers - 3)

620 -

619 def forward(

620 self,

621 input_ids: torch.Tensor | None,

vllm/model_executor/models/mistral3.py

+7 -8

@@ -44,6 +44,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

44

45 from .interfaces import (

46 MultiModalEmbeddings,

47 + SupportsEagle,

48 SupportsEagle3,

49 SupportsLoRA,

50 SupportsMultiModal,

@@ -409,7 +410,12 @@ def init_vision_tower_for_llava(

410 dummy_inputs=Mistral3DummyInputsBuilder,

411 )

412 class Mistral3ForConditionalGeneration(

412 - nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP, SupportsEagle3

413 + nn.Module,

414 + SupportsLoRA,

415 + SupportsMultiModal,

416 + SupportsPP,

417 + SupportsEagle,

418 + SupportsEagle3,

419 ):

420 packed_modules_mapping = {

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

@@ -433,13 +439,6 @@ class Mistral3ForConditionalGeneration(

439

440 raise ValueError("Only image modality is supported")

441

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

437 - self.get_language_model().model.aux_hidden_state_layers = layers

438 -

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

440 - num_layers = len(self.get_language_model().model.layers)

441 - return (2, num_layers // 2, num_layers - 3)

442 -

442 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

443 super().__init__()

444

vllm/model_executor/models/mllama4.py

+5 -9

@@ -798,20 +798,16 @@ class Llama4ForConditionalGeneration(

798 self.num_moe_layers = len(self.moe_layers)

799

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

801 - """Set which layers should output auxiliary hidden states for EAGLE3."""

801 # Delegate to underlying language model (Llama4ForCausalLM)

802 assert hasattr(self.language_model, "set_aux_hidden_state_layers")

803 self.language_model.set_aux_hidden_state_layers(layers)

804

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

807 - """Get the layer indices for auxiliary hidden state outputs.

808 -

809 - Note: The GPU model runner will override this with layers from

810 - the speculative config if available, providing dynamic configuration.

811 - """

805 + def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:

806 # Delegate to underlying language model (Llama4ForCausalLM)

813 - assert hasattr(self.language_model, "get_eagle3_aux_hidden_state_layers")

814 - return self.language_model.get_eagle3_aux_hidden_state_layers()

807 + assert hasattr(

808 + self.language_model, "get_eagle3_default_aux_hidden_state_layers"

809 + )

810 + return self.language_model.get_eagle3_default_aux_hidden_state_layers()

811

812 def set_eplb_state(

813 self,

vllm/model_executor/models/qwen2.py

+15 -15

@@ -62,7 +62,13 @@ from vllm.sequence import IntermediateTensors

62 from vllm.transformers_utils.config import is_interleaved, set_default_rope_theta

63 from vllm.v1.attention.backend import AttentionType

64

65 -from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP

65 +from .interfaces import (

66 + EagleModelMixin,

67 + SupportsEagle,

68 + SupportsEagle3,

69 + SupportsLoRA,

70 + SupportsPP,

71 +)

72 from .utils import (

73 AutoWeightsLoader,

74 PPMissingLayer,

@@ -349,7 +355,7 @@ def qwen_2_model_invariants(

355 },

356 shape_invariants=qwen_2_model_invariants,

357 )

352 -class Qwen2Model(nn.Module):

358 +class Qwen2Model(nn.Module, EagleModelMixin):

359 def __init__(

360 self,

361 *,

@@ -410,8 +416,6 @@ class Qwen2Model(nn.Module):

416 else:

417 self.norm = PPMissingLayer()

418

413 - self.aux_hidden_state_layers = tuple[int, ...]()

414 -

419 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

420 return self.embed_tokens(input_ids)

421

@@ -433,13 +437,14 @@ class Qwen2Model(nn.Module):

437 hidden_states = intermediate_tensors["hidden_states"]

438 residual = intermediate_tensors["residual"]

439

436 - aux_hidden_states = []

440 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

441 for idx, layer in enumerate(

442 islice(self.layers, self.start_layer, self.end_layer)

443 ):

440 - if idx in self.aux_hidden_state_layers:

441 - aux_hidden_states.append(hidden_states + residual)

444 hidden_states, residual = layer(positions, hidden_states, residual)

445 + self._maybe_add_hidden_state(

446 + aux_hidden_states, idx + 1, hidden_states, residual

447 + )

448

449 if not get_pp_group().is_last_rank:

450 return IntermediateTensors(

@@ -519,7 +524,9 @@ class Qwen2Model(nn.Module):

524 return loaded_params

525

526

522 -class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

527 +class Qwen2ForCausalLM(

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

529 +):

530 packed_modules_mapping = {

531 "qkv_proj": [

532 "q_proj",

@@ -566,13 +573,6 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

573 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

574 return self.model.embed_input_ids(input_ids)

575

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

570 - self.model.aux_hidden_state_layers = layers

571 -

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

573 - num_layers = len(self.model.layers)

574 - return (2, num_layers // 2, num_layers - 3)

575 -

576 def forward(

577 self,

578 input_ids: torch.Tensor | None,

vllm/model_executor/models/qwen2_5_vl.py

+2 -7

@@ -89,6 +89,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum

89

90 from .interfaces import (

91 MultiModalEmbeddings,

92 + SupportsEagle,

93 SupportsEagle3,

94 SupportsLoRA,

95 SupportsMRoPE,

@@ -1000,6 +1001,7 @@ class Qwen2_5_VLForConditionalGeneration(

1001 SupportsLoRA,

1002 SupportsPP,

1003 SupportsQuant,

1004 + SupportsEagle,

1005 SupportsEagle3,

1006 SupportsMultiModalPruning,

1007 SupportsMRoPE,

@@ -1143,13 +1145,6 @@ class Qwen2_5_VLForConditionalGeneration(

1145 self.language_model.make_empty_intermediate_tensors

1146 )

1147

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

1147 - self.language_model.model.aux_hidden_state_layers = layers

1148 -

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

1150 - num_layers = len(self.language_model.model.layers)

1151 - return (2, num_layers // 2, num_layers - 3)

1152 -

1148 def _parse_and_validate_image_input(

1149 self, **kwargs: object

1150 ) -> Qwen2_5_VLImageInputs | None:

vllm/model_executor/models/qwen3.py

+4 -9

@@ -48,7 +48,7 @@ from vllm.sequence import IntermediateTensors

48 from vllm.transformers_utils.config import set_default_rope_theta

49 from vllm.v1.attention.backend import AttentionType

50

51 -from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP

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

52 from .qwen2 import Qwen2MLP as Qwen3MLP

53 from .qwen2 import Qwen2Model

54 from .utils import AutoWeightsLoader, PPMissingLayer, extract_layer_index, maybe_prefix

@@ -258,7 +258,9 @@ class Qwen3Model(Qwen2Model):

258 )

259

260

261 -class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

261 +class Qwen3ForCausalLM(

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

263 +):

264 packed_modules_mapping = {

265 "qkv_proj": [

266 "q_proj",

@@ -307,13 +309,6 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):

309 self.model.make_empty_intermediate_tensors

310 )

311

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

311 - self.model.aux_hidden_state_layers = layers

312 -

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

314 - num_layers = len(self.model.layers)

315 - return (2, num_layers // 2, num_layers - 3)

316 -

312 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

313 return self.model.embed_input_ids(input_ids)

314

vllm/model_executor/models/qwen3_moe.py

+16 -19

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

65 from vllm.model_executor.models.utils import sequence_parallel_chunk

66 from vllm.sequence import IntermediateTensors

67

68 -from .interfaces import MixtureOfExperts, SupportsEagle3, SupportsLoRA, SupportsPP

68 +from .interfaces import (

69 + EagleModelMixin,

70 + MixtureOfExperts,

71 + SupportsEagle,

72 + SupportsEagle3,

73 + SupportsLoRA,

74 + SupportsPP,

75 +)

76 from .utils import (

77 AutoWeightsLoader,

78 PPMissingLayer,

@@ -427,7 +434,7 @@ class Qwen3MoeDecoderLayer(nn.Module):

434

435

436 @support_torch_compile

430 -class Qwen3MoeModel(nn.Module):

437 +class Qwen3MoeModel(nn.Module, EagleModelMixin):

438 def __init__(

439 self,

440 *,

@@ -461,8 +468,6 @@ class Qwen3MoeModel(nn.Module):

468 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

469 ["hidden_states", "residual"], config.hidden_size

470 )

464 - # Track layers for auxiliary hidden state outputs (EAGLE3)

465 - self.aux_hidden_state_layers: tuple[int, ...] = ()

471

472 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

473 return self.embed_tokens(input_ids)

@@ -485,18 +490,17 @@ class Qwen3MoeModel(nn.Module):

490 hidden_states = intermediate_tensors["hidden_states"]

491 residual = intermediate_tensors["residual"]

492

488 - aux_hidden_states = []

493 + aux_hidden_states = self._maybe_add_hidden_state(

494 + [], self.start_layer, hidden_states, residual

495 + )

496 for layer_idx, layer in enumerate(

497 islice(self.layers, self.start_layer, self.end_layer),

498 start=self.start_layer,

499 ):

493 - # Collect auxiliary hidden states if specified

494 - if layer_idx in self.aux_hidden_state_layers:

495 - aux_hidden_state = (

496 - hidden_states + residual if residual is not None else hidden_states

497 - )

498 - aux_hidden_states.append(aux_hidden_state)

500 hidden_states, residual = layer(positions, hidden_states, residual)

501 + self._maybe_add_hidden_state(

502 + aux_hidden_states, layer_idx + 1, hidden_states, residual

503 + )

504

505 if not get_pp_group().is_last_rank:

506 return IntermediateTensors(

@@ -666,7 +670,7 @@ class Qwen3MoeModel(nn.Module):

670

671

672 class Qwen3MoeForCausalLM(

669 - nn.Module, SupportsPP, SupportsLoRA, SupportsEagle3, MixtureOfExperts

673 + nn.Module, SupportsPP, SupportsLoRA, SupportsEagle, SupportsEagle3, MixtureOfExperts

674 ):

675 packed_modules_mapping = {

676 "qkv_proj": [

@@ -751,13 +755,6 @@ class Qwen3MoeForCausalLM(

755 moe.n_redundant_experts = self.num_redundant_experts

756 moe.experts.update_expert_map()

757

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

755 - self.model.aux_hidden_state_layers = layers

756 -

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

758 - num_layers = len(self.model.layers)

759 - return (2, num_layers // 2, num_layers - 3)

760 -

758 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

759 return self.model.embed_input_ids(input_ids)

760

vllm/model_executor/models/qwen3_vl.py

+6 -11

@@ -101,6 +101,7 @@ from vllm.utils.math_utils import round_up

101

102 from .interfaces import (

103 MultiModalEmbeddings,

104 + SupportsEagle,

105 SupportsEagle3,

106 SupportsLoRA,

107 SupportsMRoPE,

@@ -1275,13 +1276,10 @@ class Qwen3LLMModel(Qwen3Model):

1276 hidden_states = intermediate_tensors["hidden_states"]

1277 residual = intermediate_tensors["residual"]

1278

1278 - aux_hidden_states = []

1279 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

1280 for layer_idx, layer in islice(

1281 enumerate(self.layers), self.start_layer, self.end_layer

1282 ):

1282 - if layer_idx in self.aux_hidden_state_layers:

1283 - aux_hidden_states.append(hidden_states + residual)

1284 -

1283 hidden_states, residual = layer(

1284 positions,

1285 hidden_states,

@@ -1295,6 +1293,9 @@ class Qwen3LLMModel(Qwen3Model):

1293 hidden_states

1294 + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]

1295 )

1296 + self._maybe_add_hidden_state(

1297 + aux_hidden_states, layer_idx + 1, hidden_states, residual

1298 + )

1299

1300 if not get_pp_group().is_last_rank:

1301 return IntermediateTensors(

@@ -1351,6 +1352,7 @@ class Qwen3VLForConditionalGeneration(

1352 SupportsLoRA,

1353 SupportsPP,

1354 SupportsMRoPE,

1355 + SupportsEagle,

1356 SupportsEagle3,

1357 SupportsMultiModalPruning,

1358 ):

@@ -1449,13 +1451,6 @@ class Qwen3VLForConditionalGeneration(

1451 self.language_model.make_empty_intermediate_tensors

1452 )

1453

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

1453 - self.language_model.model.aux_hidden_state_layers = layers

1454 -

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

1456 - num_layers = len(self.language_model.model.layers)

1457 - return (2, num_layers // 2, num_layers - 3)

1458 -

1454 def _get_deepstack_input_embeds(

1455 self,

1456 num_tokens: int,

vllm/model_executor/models/qwen3_vl_moe.py

+7 -5

@@ -102,19 +102,17 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):

102 hidden_states = intermediate_tensors["hidden_states"]

103 residual = intermediate_tensors["residual"]

104

105 - aux_hidden_states = []

105 + aux_hidden_states = self._maybe_add_hidden_state(

106 + [], self.start_layer, hidden_states, residual

107 + )

108 for layer_idx, layer in islice(

109 enumerate(self.layers), self.start_layer, self.end_layer

110 ):

109 - if layer_idx in self.aux_hidden_state_layers:

110 - aux_hidden_states.append(hidden_states + residual)

111 -

111 hidden_states, residual = layer(

112 positions,

113 hidden_states,

114 residual,

115 )

117 -

116 if deepstack_input_embeds is not None and layer_idx in range(

117 0, len(deepstack_input_embeds)

118 ):

@@ -123,6 +121,10 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):

121 + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]

122 )

123

124 + self._maybe_add_hidden_state(

125 + aux_hidden_states, layer_idx + 1, hidden_states, residual

126 + )

127 +

128 if not get_pp_group().is_last_rank:

129 return IntermediateTensors(

130 {"hidden_states": hidden_states, "residual": residual}

vllm/model_executor/models/step1.py

+12 -12

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

31 VocabParallelEmbedding,

32 )

33 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

34 -from vllm.model_executor.models.interfaces import SupportsPP

34 +from vllm.model_executor.models.interfaces import (

35 + EagleModelMixin,

36 + SupportsEagle,

37 + SupportsEagle3,

38 + SupportsPP,

39 +)

40 from vllm.model_executor.models.utils import (

41 AutoWeightsLoader,

42 PPMissingLayer,

@@ -274,7 +279,7 @@ class StepDecoderLayer(nn.Module):

279 return loaded_params

280

281

277 -class StepDecoderModel(nn.Module):

282 +class StepDecoderModel(nn.Module, EagleModelMixin):

283 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

284 super().__init__()

285 config = vllm_config.model_config.hf_config

@@ -303,9 +308,6 @@ class StepDecoderModel(nn.Module):

308 else:

309 self.norm = PPMissingLayer()

310

306 - self.aux_hidden_state_layers: tuple[int, ...] = getattr(

307 - config, "aux_hidden_state_layers", ()

308 - )

311 self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

312 ["hidden_states", "residual"],

313 config.hidden_size,

@@ -333,14 +335,12 @@ class StepDecoderModel(nn.Module):

335 hidden_states = intermediate_tensors["hidden_states"]

336 residual = intermediate_tensors["residual"]

337

336 - aux_hidden_states = []

338 + aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

339 for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):

338 - if idx in self.aux_hidden_state_layers:

339 - if residual is None:

340 - aux_hidden_states.append(hidden_states)

341 - else:

342 - aux_hidden_states.append(hidden_states + residual)

340 hidden_states, residual = layer(positions, hidden_states, residual)

341 + self._maybe_add_hidden_state(

342 + aux_hidden_states, idx + 1, hidden_states, residual

343 + )

344

345 if not get_pp_group().is_last_rank:

346 return IntermediateTensors(

@@ -353,7 +353,7 @@ class StepDecoderModel(nn.Module):

353 return hidden_states

354

355

356 -class Step1ForCausalLM(nn.Module, SupportsPP):

356 +class Step1ForCausalLM(nn.Module, SupportsPP, SupportsEagle, SupportsEagle3):

357 packed_modules_mapping = STEP_PACKED_MODULES_MAPPING

358

359 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

vllm/model_executor/models/transformers/base.py

+1 -1

@@ -618,6 +618,6 @@ class Base(

618 # Ensure that the capture hooks are installed before dynamo traces the model

619 maybe_install_capturing_hooks(self.model)

620

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

621 + def get_eagle3_default_aux_hidden_state_layers(self) -> tuple[int, ...]:

622 num_layers = self.text_config.num_hidden_layers

623 return (2, num_layers // 2, num_layers - 3)

vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py

+1 -1

@@ -27,7 +27,7 @@ def set_eagle3_aux_hidden_state_layers(

27 if aux_layers:

28 logger.info("Using Eagle3 auxiliary layers from config: %s", aux_layers)

29 else:

30 - aux_layers = eagle3_model.get_eagle3_aux_hidden_state_layers()

30 + aux_layers = eagle3_model.get_eagle3_default_aux_hidden_state_layers()

31 logger.info("Using Eagle3 auxiliary layers from model: %s", aux_layers)

32 eagle3_model.set_aux_hidden_state_layers(aux_layers)

33

vllm/v1/worker/gpu_model_runner.py

+3 -1

@@ -4556,7 +4556,9 @@ class GPUModelRunner(

4556 aux_layers,

4557 )

4558 else:

4559 - aux_layers = self.model.get_eagle3_aux_hidden_state_layers()

4559 + aux_layers = (

4560 + self.model.get_eagle3_default_aux_hidden_state_layers()

4561 + )

4562

4563 self.model.set_aux_hidden_state_layers(aux_layers)

4564 time_after_load = time.perf_counter()

[Refactor] Consolidate SupportsEagle (#36063)