[NPU] bugfix with Kimi-k2 and bge-reranker-v2 model (#17478)

Co-authored-by: amote-i <49533125+amote-i@users.noreply.github.com>
Co-authored-by: cy <chenyang08056032@163.com>

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -675,7 +675,14 @@ class AscendAttnBackend(AttentionBackend):
                 )
 
             else:
-                if layer.qk_head_dim <= 128:
+                causal = True
+                if (
+                    layer.is_cross_attention
+                    or layer.attn_type == AttentionType.ENCODER_ONLY
+                ):
+                    causal = False
+
+                if layer.qk_head_dim <= 128 and causal:
                     query = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
                     attn_output = torch.empty(
                         (query.shape[0], layer.tp_q_head_num * layer.v_head_dim),
@@ -709,13 +716,6 @@ class AscendAttnBackend(AttentionBackend):
                     q_ = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim)
                     o_ = attn_output.view(-1, layer.tp_q_head_num, layer.v_head_dim)
 
-                    causal = True
-                    if (
-                        layer.is_cross_attention
-                        or layer.attn_type == AttentionType.ENCODER_ONLY
-                    ):
-                        causal = False
-
                     self.native_attn._run_sdpa_forward_extend(
                         q_,
                         o_,

python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -1648,7 +1648,7 @@ class NPUCompressedTensorsW4A16Int4DynamicMoEMethod(CompressedTensorsMoEMethod):
         self.num_experts = num_experts
         if (
             extra_weight_attrs.get(
-                "intermediate_size_full", intermediate_size_per_partition
+                "moe_intermediate_size", intermediate_size_per_partition
             )
             // intermediate_size_per_partition
             > 1

python/sglang/srt/layers/rotary_embedding.py

@@ -115,9 +115,10 @@ class RotaryEmbedding(MultiPlatformOp):
             cache = cache.to(dtype)
 
         if (
-            (not (_is_cuda or _is_npu) or self.head_size not in [64, 128, 256, 512])
+            (not (_is_cuda) or self.head_size not in [64, 128, 256, 512])
             and not (_is_cpu)
             and not (_is_xpu)
+            and not (_is_npu)
         ):
             if _is_cuda or _is_hip:
                 from sgl_kernel import rotary_embedding

python/sglang/srt/models/bert.py

@@ -17,6 +17,7 @@ from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import add_prefix
 
 BertConfig = None
@@ -365,10 +366,15 @@ class BertModel(nn.Module):
             quant_config=quant_config,
             prefix=add_prefix("encoder", prefix),
         )
+        pooling_type = (
+            PoolingType.CLS
+            if get_global_server_args().is_embedding
+            else PoolingType.LAST
+        )
         self.pooler = (
             BertPooler(config)
             if self.use_bert_pooler
-            else Pooler(pooling_type=PoolingType.LAST, normalize=True)
+            else Pooler(pooling_type=pooling_type, normalize=True)
         )
 
     @torch.no_grad()