Unverified Commit 66face35 authored by YanbingJiang's avatar YanbingJiang Committed by GitHub
Browse files

Fix CI failure of TypeError: RotaryEmbedding.forward_cpu() got an unexpected...

Fix CI failure of TypeError: RotaryEmbedding.forward_cpu() got an unexpected keyword argument 'fused_set_kv_buffer_arg' (#11009)
parent 5519766a
...@@ -206,6 +206,7 @@ class RotaryEmbedding(CustomOp): ...@@ -206,6 +206,7 @@ class RotaryEmbedding(CustomOp):
query: torch.Tensor, query: torch.Tensor,
key: torch.Tensor, key: torch.Tensor,
offsets: Optional[torch.Tensor] = None, offsets: Optional[torch.Tensor] = None,
fused_set_kv_buffer_arg=None,
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
positions = torch.add(positions, offsets) if offsets is not None else positions positions = torch.add(positions, offsets) if offsets is not None else positions
if _is_cpu_amx_available: if _is_cpu_amx_available:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment