fix the crash of meta-llama/Llama-3.2-1B (#2918)

* fix the crash of meta-llama/Llama-3.2-1B

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>

* Apply suggestions from code review

Simpler fix (which doesn't break vlms).

---------

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
Co-authored-by: Nicolas Patry <patry.nicolas@protonmail.com>

Author: sywangyi

server/text_generation_server/models/custom_modeling/flash_llama_modeling.py

@@ -642,9 +642,7 @@ def __init__(self, prefix: str, config, weights, name=None):
         embedding_multiplier = getattr(config, "embedding_multiplier", None)
         if embedding_multiplier is not None:
             self.embed_tokens.weight.data *= embedding_multiplier
-
-        prefix = "lm_head" if not prefix or name != "model" else f"{prefix}.{suffix}"
-
+        prefix = suffix if not prefix or name != "model" else f"{prefix}.{suffix}"
         with no_fp8(weights):
             self.lm_head = SpeculativeHead.load(
                 config,