llama4 npu adapt (#17123)

Co-authored-by: cy <chenyang08056032@163.com>
Co-authored-by: sglang-npu-bot <sglangnpu@163.com>

python/sglang/srt/models/llama4.py

@@ -56,11 +56,13 @@ from sglang.srt.utils import (
     fast_topk,
     get_compiler_backend,
     is_cuda,
+    is_npu,
     make_layers,
 )
 from sglang.srt.utils.common import get_current_device_stream_fast
 
 _is_cuda = is_cuda()
+_is_npu = is_npu()
 
 logger = logging.getLogger(__name__)
 
@@ -329,6 +331,8 @@ class Llama4Attention(nn.Module):
         if self.rotary_emb is not None:
             q_view, k_view = qk.split([self.q_size, self.kv_size], dim=-1)
             q_out_unused, k_out_unused = self.rotary_emb(positions, q_view, k_view)
+            if _is_npu:
+                qk = torch.cat([q_out_unused, k_out_unused], dim=-1)
             del q_view, k_view, q_out_unused, k_out_unused
 
         if self.qk_norm is not None:

python/sglang/srt/server_args.py

@@ -1480,9 +1480,10 @@ class ServerArgs:
                 "fa3",
                 "aiter",
                 "triton",
+                "ascend",
                 "trtllm_mha",
                 "intel_xpu",
-            }, f"fa3, aiter, triton, trtllm_mha or intel_xpu is required for Llama4 model but got {self.attention_backend}"
+            }, f"fa3, aiter, triton, ascend, trtllm_mha or intel_xpu is required for Llama4 model but got {self.attention_backend}"
             if is_sm100_supported() and self.moe_runner_backend == "auto":
                 if self.quantization in {"fp8", "modelopt_fp8"}:
                     self.moe_runner_backend = "flashinfer_trtllm"