[gaudi] Fix the Llama-4-Maverick-17B-128E crash issue (#3246)

yuanwu2017 · web-flow · commit 6b6e30a6f680 · 2025-05-29T11:38:44.000+02:00
Signed-off-by: yuanwu &lt;yuan.wu@intel.com&gt;
diff --git a/backends/gaudi/server/text_generation_server/models/custom_modeling/flash_llama4_modeling.py b/backends/gaudi/server/text_generation_server/models/custom_modeling/flash_llama4_modeling.py
@@ -48,7 +48,6 @@
 )
 from text_generation_server.models.custom_modeling.flash_llama_modeling import (
     FlashLlamaAttention,
-    LlamaMLP,
 )
 
 
@@ -444,7 +443,7 @@ def __init__(self, prefix, config, weights, layer_idx):
         if self.is_moe_layer:  # the 128E model interleaves dense / sparse
             self.feed_forward = Llama4TextMoe(f"{prefix}.feed_forward", config, weights)
         else:
-            self.feed_forward = LlamaMLP(f"{prefix}.feed_forward", config, weights)
+            self.feed_forward = Llama4TextMLP(f"{prefix}.feed_forward", config, weights)
 
         self.input_layernorm = FastRMSNorm.load(
             prefix=f"{prefix}.input_layernorm",

Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,6 @@`
`48`	`48`	`)`
`49`	`49`	`from text_generation_server.models.custom_modeling.flash_llama_modeling import (`
`50`	`50`	`FlashLlamaAttention,`
`51`		`- LlamaMLP,`
`52`	`51`	`)`
`53`	`52`
`54`	`53`
`@@ -444,7 +443,7 @@ def __init__(self, prefix, config, weights, layer_idx):`
`444`	`443`	`if self.is_moe_layer: # the 128E model interleaves dense / sparse`
`445`	`444`	`self.feed_forward = Llama4TextMoe(f"{prefix}.feed_forward", config, weights)`
`446`	`445`	`else:`
`447`		`- self.feed_forward = LlamaMLP(f"{prefix}.feed_forward", config, weights)`
	`446`	`+ self.feed_forward = Llama4TextMLP(f"{prefix}.feed_forward", config, weights)`
`448`	`447`
`449`	`448`	`self.input_layernorm = FastRMSNorm.load(`
`450`	`449`	`prefix=f"{prefix}.input_layernorm",`