[Bugfix] Fix MTP accuracy for GLM-5 (#34385)

Signed-off-by: mgoin <mgoin64@gmail.com>
(cherry picked from commit ec12d39d44)

benchmarks/kernels/benchmark_moe.py

@@ -686,6 +686,7 @@ def get_model_params(config):
         "DeepseekV2ForCausalLM",
         "DeepseekV3ForCausalLM",
         "DeepseekV32ForCausalLM",
+        "GlmMoeDsaForCausalLM",
         "Glm4MoeForCausalLM",
         "Glm4MoeLiteForCausalLM",
         "NemotronHForCausalLM",

requirements/nightly_torch_test.txt

@@ -43,5 +43,5 @@ tritonclient>=2.51.0
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.1.10
+fastsafetensors>=0.2.2
 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/test.in

@@ -53,7 +53,7 @@ arctic-inference == 0.1.1 # Required for suffix decoding test
 numba == 0.61.2 # Required for N-gram speculative decoding
 numpy
 runai-model-streamer[s3,gcs]==0.15.3
-fastsafetensors>=0.1.10
+fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
 pydantic>=2.12 # 2.11 leads to error on python 3.13
 decord==0.6.0
 terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test

requirements/test.txt

@@ -224,7 +224,7 @@ fastparquet==2024.11.0
     # via genai-perf
 fastrlock==0.8.2
     # via cupy-cuda12x
-fastsafetensors==0.1.10
+fastsafetensors==0.2.2
     # via -r requirements/test.in
 filelock==3.16.1
     # via
@@ -1174,7 +1174,6 @@ torch==2.10.0+cu129
     #   bitsandbytes
     #   efficientnet-pytorch
     #   encodec
-    #   fastsafetensors
     #   kornia
     #   lightly
     #   lightning

setup.py

@@ -1035,7 +1035,7 @@ setup(
     extras_require={
         "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy"],
         "tensorizer": ["tensorizer==2.10.1"],
-        "fastsafetensors": ["fastsafetensors >= 0.1.10"],
+        "fastsafetensors": ["fastsafetensors >= 0.2.2"],
         "runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
         "audio": [
             "librosa",

tests/entrypoints/openai/test_realtime_validation.py

@@ -129,5 +129,5 @@ async def test_multi_chunk_streaming(
                 " First words I spoke in the original phonograph."
                 " A little piece of practical poetry. Mary had a little lamb,"
                 " it sleeps with quite a flow, and everywhere that Mary went,"
-                " the lamb was sure to go"
+                " the lamb was sure to go."
             )

tests/models/registry.py

@@ -275,6 +275,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
         "zai-org/GLM-4.7-Flash",
         min_transformers_version="5.0.0",
     ),
+    "GlmMoeDsaForCausalLM": _HfExamplesInfo(
+        "zai-org/GLM-5", min_transformers_version="5.0.1", is_available_online=False
+    ),
     "GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}),
     "GPTBigCodeForCausalLM": _HfExamplesInfo(
         "bigcode/starcoder",

tests/models/test_initialization.py

@@ -97,7 +97,7 @@ def can_initialize(
             "pickle error when loading `transformers.models.auto.CONFIG_MAPPING`"
         )
 
-    if model_arch == "DeepseekV32ForCausalLM":
+    if model_arch in ["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"]:
         from vllm.platforms import current_platform
 
         capability = current_platform.get_device_capability()

vllm/config/speculative.py

@@ -181,7 +181,7 @@ class SpeculativeConfig:
     @staticmethod
     def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
         initial_architecture = hf_config.architectures[0]
-        if hf_config.model_type in ("deepseek_v3", "deepseek_v32"):
+        if hf_config.model_type in ("deepseek_v3", "deepseek_v32", "glm_moe_dsa"):
             hf_config.model_type = "deepseek_mtp"
         if hf_config.model_type == "deepseek_mtp":
             n_predict = getattr(hf_config, "num_nextn_predict_layers", None)

vllm/entrypoints/openai/realtime/connection.py

@@ -48,7 +48,6 @@ class RealtimeConnection:
         self.generation_task: asyncio.Task | None = None
 
         self._is_connected = False
-        self._is_input_finished = False
         self._is_model_validated = False
 
         self._max_audio_filesize_mb = envs.VLLM_MAX_AUDIO_CLIP_FILESIZE_MB
@@ -145,7 +144,7 @@ class RealtimeConnection:
             commit_event = InputAudioBufferCommit(**event)
             # final signals that the audio is finished
             if commit_event.final:
-                self._is_input_finished = True
+                self.audio_queue.put_nowait(None)
             else:
                 await self.start_generation()
         else:
@@ -239,11 +238,6 @@ class RealtimeConnection:
                     # finish because websocket connection was killed
                     break
 
-                if self.audio_queue.empty() and self._is_input_finished:
-                    # finish because client signals that audio input
-                    # is finished
-                    break
-
             usage = UsageInfo(
                 prompt_tokens=prompt_token_ids_len,
                 completion_tokens=completion_tokens_len,

vllm/model_executor/models/deepseek_v2.py

@@ -836,7 +836,7 @@ class DeepseekV2MLAAttention(nn.Module):
                 qk_rope_head_dim,
                 max_position=max_position_embeddings,
                 rope_parameters=config.rope_parameters,
-                is_neox_style=True,
+                is_neox_style=not getattr(config, "indexer_rope_interleave", False),
             )
             self.indexer = Indexer(
                 vllm_config,
@@ -1499,6 +1499,10 @@ class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM):
     pass
 
 
+class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM):
+    pass
+
+
 # Compatibility with
 # https://huggingface.co/deepseek-ai/DeepSeek-V3-Base/blob/main/configuration_deepseek.py
 def get_spec_layer_idx_from_weight_name(

vllm/model_executor/models/registry.py

@@ -114,6 +114,7 @@ _TEXT_GENERATION_MODELS = {
     "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),
     "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),
     "Glm4MoeLiteForCausalLM": ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
+    "GlmMoeDsaForCausalLM": ("deepseek_v2", "GlmMoeDsaForCausalLM"),
     "GptOssForCausalLM": ("gpt_oss", "GptOssForCausalLM"),
     "GPT2LMHeadModel": ("gpt2", "GPT2LMHeadModel"),
     "GPTBigCodeForCausalLM": ("gpt_bigcode", "GPTBigCodeForCausalLM"),

vllm/transformers_utils/model_arch_config_convertor.py

@@ -237,6 +237,7 @@ class ModelArchConfigConvertorBase:
             "deepseek_v3",
             "deepseek_v32",
             "deepseek_mtp",
+            "glm_moe_dsa",
             "glm4_moe_lite",
             "glm4_moe_lite_mtp",
             "kimi_k2",

vllm/v1/spec_decode/eagle.py

@@ -1503,6 +1503,24 @@ class SpecDecodeBaseProposer:
                 del self.model.lm_head
             self.model.lm_head = target_language_model.lm_head
 
+            # MTP models call compute_logits via shared_head.head (a
+            # ParallelLMHead inside each MTP layer), not self.model.lm_head.
+            # If the checkpoint omits a copy of the lm_head weights at the
+            # MTP layer path, shared_head.head stays uninitialised and
+            # produces NaN logits. Always share it explicitly.
+            inner = getattr(self.model, "model", None)
+            layers = getattr(inner, "layers", None) if inner else None
+            if layers is not None:
+                items = layers.values() if isinstance(layers, nn.ModuleDict) else layers
+                for layer in items:
+                    sh = getattr(layer, "shared_head", None)
+                    if sh is not None and hasattr(sh, "head"):
+                        del sh.head
+                        sh.head = target_language_model.lm_head
+                        logger.info(
+                            "Shared target model lm_head with MTP shared_head.head."
+                        )
+
     @torch.inference_mode()
     def dummy_run(
         self,