Fix CI failure of TypeError: RotaryEmbedding.forward_cpu() got an unexpected keyword argument 'fused_set_kv_buffer_arg' (#11009)
This commit is contained in:
@@ -206,6 +206,7 @@ class RotaryEmbedding(CustomOp):
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
offsets: Optional[torch.Tensor] = None,
|
||||
fused_set_kv_buffer_arg=None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
positions = torch.add(positions, offsets) if offsets is not None else positions
|
||||
if _is_cpu_amx_available:
|
||||
|
||||
Reference in New Issue
Block a user