[Feature] npu support enable_torch_compile for torchair backend (#13410)
Co-authored-by: ZhengdQin <zhengdqin@gmail.com>
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user