[Feature] npu support enable_torch_compile for torchair backend (#13410)

Co-authored-by: ZhengdQin <zhengdqin@gmail.com>
This commit is contained in:
XDaoHong
2025-12-16 09:23:51 +08:00
committed by GitHub
co-authored by ZhengdQin
parent 3ffa260474
commit 4733fcff1f
7 changed files with 243 additions and 50 deletions
+15 -20
View File
@@ -269,27 +269,22 @@ class RotaryEmbedding(CustomOp):
fused_set_kv_buffer_arg is None
), "fused_set_kv_buffer_arg is not supported for npu implementation"
if get_bool_env_var("SGLANG_ENABLE_TORCH_COMPILE"):
return self.forward_native(
positions, query, key, offsets, fused_set_kv_buffer_arg
)
else:
rotary_mode = "half"
if self.is_neox_style:
rotary_mode = "half"
if self.is_neox_style:
rotary_mode = "half"
else:
rotary_mode = "interleave"
mrope_section = [0, 0, 0]
query_out, key_out = torch_npu.npu_mrope(
positions,
query,
key,
self.cos_sin_cache,
self.head_size,
mrope_section=mrope_section,
rotary_mode=rotary_mode,
)
return query_out, key_out
else:
rotary_mode = "interleave"
mrope_section = [0, 0, 0]
query_out, key_out = torch_npu.npu_mrope(
positions,
query,
key,
self.cos_sin_cache,
self.head_size,
mrope_section=mrope_section,
rotary_mode=rotary_mode,
)
return query_out, key_out
def forward_cpu(
self,