From a245cae3d1e7a5bae6028fc9f5da32cda507a6cb Mon Sep 17 00:00:00 2001
From: ZeyuanChen2000 <806126546@qq.com>
Date: Thu, 28 May 2026 16:05:39 +0800
Subject: [PATCH] [NPU] fix model ERNIE-4.5-21B-A3B-PT bias need 1D error
 (#26038)

---
 python/sglang/srt/models/ernie4.py | 10 ++++++++--
 python/sglang/srt/models/llama.py  |  1 +
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/python/sglang/srt/models/ernie4.py b/python/sglang/srt/models/ernie4.py
index 2d33874da..7c292a125 100644
--- a/python/sglang/srt/models/ernie4.py
+++ b/python/sglang/srt/models/ernie4.py
@@ -42,9 +42,11 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4MLP
 from sglang.srt.models.llama import LlamaAttention as Ernie4Attention
-from sglang.srt.utils import add_prefix, make_layers
+from sglang.srt.utils import add_prefix, is_npu, make_layers
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
+_is_npu = is_npu()
+
 
 class MoEGate(nn.Module):
     def __init__(
@@ -86,12 +88,16 @@ class Ernie4Moe(nn.Module):
 
         self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix))
 
+        correction_bias = self.gate.e_score_correction_bias
+        # npu only supports 1D, but current correction_bias is 2D
+        if _is_npu:
+            correction_bias = correction_bias.squeeze(0)
         self.topk = TopK(
             top_k=config.moe_k,
             layer_id=layer_id,
             renormalize=True,
             use_grouped_topk=False,
-            correction_bias=self.gate.e_score_correction_bias,
+            correction_bias=correction_bias,
         )
 
         self.experts = get_moe_impl_class(quant_config)(
diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py
index 22fa71563..742f810ac 100644
--- a/python/sglang/srt/models/llama.py
+++ b/python/sglang/srt/models/llama.py
@@ -222,6 +222,7 @@ class LlamaAttention(nn.Module):
             self.q_size,
             self.kv_size,
             self.head_dim,
+            is_neox_style=self.rotary_emb.is_neox_style,
         )
         return q, k, v