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 14, 05:30 UTC
[Bug] [Spec Decode] Fix model_initialization test and mismatch in aux_hidden_layers (#24613)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com> Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Wenlong Wang committed
Sep 10, 2025 at 21:23 UTC
6c8deacd7211ce4f536d22938e26c9ab931bb1f0
3 files changed
+32
-11
tests/models/registry.py
+16
-7
@@ -97,6 +97,12 @@ class _HfExamplesInfo:
97
max_num_seqs: Optional[int] = None
98
"""Maximum number of sequences to be processed in a single iteration."""
99
100
+ use_original_num_layers: bool = False
101
+ """
102
+ If True, use the original number of layers from the model config
103
+ instead of minimal layers for testing.
104
+ """
105
+
106
def check_transformers_version(
107
self,
108
*,
@@ -597,18 +603,21 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
603
"EagleDeepSeekMTPModel": _HfExamplesInfo("eagle618/deepseek-v3-random",
604
speculative_model="eagle618/eagle-deepseek-v3-random", # noqa: E501
605
trust_remote_code=True),
600
- "EagleLlamaForCausalLM": _HfExamplesInfo("yuhuili/EAGLE-LLaMA3-Instruct-8B",
606
+ "EagleLlamaForCausalLM": _HfExamplesInfo("meta-llama/Meta-Llama-3-8B-Instruct", # noqa: E501
607
trust_remote_code=True,
608
speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
603
- tokenizer="meta-llama/Meta-Llama-3-8B-Instruct"), # noqa: E501
604
- "Eagle3LlamaForCausalLM": _HfExamplesInfo("yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", # noqa: E501
609
+ tokenizer="meta-llama/Meta-Llama-3-8B-Instruct"), # noqa: E501
610
+ "Eagle3LlamaForCausalLM": _HfExamplesInfo("meta-llama/Llama-3.1-8B-Instruct", # noqa: E501
611
trust_remote_code=True,
606
- speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
607
- tokenizer="meta-llama/Llama-3.1-8B-Instruct"),
608
- "LlamaForCausalLMEagle3": _HfExamplesInfo("AngelSlim/Qwen3-8B_eagle3", # noqa: E501
612
+ speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", # noqa: E501
613
+ tokenizer="meta-llama/Llama-3.1-8B-Instruct",
614
+ use_original_num_layers=True,
615
+ max_model_len=10240),
616
+ "LlamaForCausalLMEagle3": _HfExamplesInfo("Qwen/Qwen3-8B", # noqa: E501
617
trust_remote_code=True,
618
speculative_model="AngelSlim/Qwen3-8B_eagle3", # noqa: E501
611
- tokenizer="Qwen/Qwen3-8B"),
619
+ tokenizer="Qwen/Qwen3-8B",
620
+ use_original_num_layers=True),
621
"EagleLlama4ForCausalLM": _HfExamplesInfo(
622
"morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
623
trust_remote_code=True,
tests/models/test_initialization.py
+4
-1
@@ -36,7 +36,10 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
36
37
hf_overrides_fn = partial(dummy_hf_overrides,
38
model_arch=model_arch,
39
- exist_overrides=model_info.hf_overrides)
39
+ exist_overrides=model_info.hf_overrides,
40
+ use_original_num_layers=getattr(
41
+ model_info, 'use_original_num_layers',
42
+ False))
43
44
# Avoid calling model.forward()
45
def _initialize_kv_caches_v0(self) -> None:
tests/models/utils.py
+12
-3
@@ -396,6 +396,7 @@ def dummy_hf_overrides(
396
*,
397
model_arch: str = "",
398
exist_overrides: Optional[dict[str, Any]] = None,
399
+ use_original_num_layers: bool = False,
400
) -> PretrainedConfig:
401
"""
402
Dummy HF overrides function used to create dummy model
@@ -412,10 +413,18 @@ def dummy_hf_overrides(
413
414
# we use three layers for Gemma-3n to check
415
# both normal layer and kv_shared_layer
415
- num_hidden_layers = (3 if model_arch == "Gemma3nForConditionalGeneration"
416
- else 1)
416
+ if use_original_num_layers:
417
+ # Use the original number of layers from the config
418
+ num_layers = getattr(text_config, 'num_layers', 1)
419
+ num_hidden_layers = getattr(text_config, 'num_hidden_layers', 1)
420
+ else:
421
+ # Use minimal layers for testing
422
+ num_layers = 1
423
+ num_hidden_layers = (3 if model_arch
424
+ == "Gemma3nForConditionalGeneration" else 1)
425
+
426
text_config.update({
418
- "num_layers": 1,
427
+ "num_layers": num_layers,
428
"num_hidden_layers": num_hidden_layers,
429
"num_experts": num_experts,
430
"num_experts_per_tok": 2,