diff --git a/tests/config/test_speculative_draft_hf_overrides.py b/tests/config/test_speculative_draft_hf_overrides.py index e7014faa47d0..39ebd7924bf2 100644 --- a/tests/config/test_speculative_draft_hf_overrides.py +++ b/tests/config/test_speculative_draft_hf_overrides.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch import pytest -from transformers import PreTrainedConfig +from transformers import Glm5NextConfig, PreTrainedConfig from vllm.config.parallel import ParallelConfig from vllm.config.speculative import SpeculativeConfig @@ -270,3 +270,21 @@ def test_mtp_index_share_override( speculative_config.draft_model_config.hf_config.index_share_for_mtp_iteration is expected ) + + +@pytest.mark.cpu_test +def test_glm5_next_mtp_override_promotes_index_share(): + """The flag lives on text_config, but the MTP proposers read the draft config.""" + config = Glm5NextConfig( + architectures=["Glm5NextForConditionalGeneration"], + text_config={ + "num_nextn_predict_layers": 1, + "index_share_for_mtp_iteration": True, + }, + ) + + draft_config = SpeculativeConfig.hf_config_override(config) + + assert draft_config.architectures == ["Glm5NextMTPModel"] + assert draft_config.n_predict == 1 + assert draft_config.index_share_for_mtp_iteration is True diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 4c61ead82dab..b916d7da2bed 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -1058,9 +1058,16 @@ def hf_config_override(hf_config: PreTrainedConfig) -> PreTrainedConfig: ) if hf_config.model_type == "glm5_next": hf_config.model_type = "glm5_next_mtp" - n_predict = hf_config.get_text_config().num_nextn_predict_layers + text_config = hf_config.get_text_config() + share_mtp_indices = getattr( + text_config, "index_share_for_mtp_iteration", False + ) hf_config.update( - {"n_predict": n_predict, "architectures": ["Glm5NextMTPModel"]} + { + "n_predict": text_config.num_nextn_predict_layers, + "architectures": ["Glm5NextMTPModel"], + "index_share_for_mtp_iteration": share_mtp_indices, + } ) return hf_config