
    .`i                         d dl Z d dlmZ d dlmZ d dlmZ d dlmZm	Z	 d dl
mZ  G d de	          Z G d	 d
e          ZdS )    N)Iterable)	BlockHash)LoadStoreSpec)BackendBlockStatusCPULoadStoreSpecc                   H     e Zd Zej        dej        fgz   Zdef fdZ xZ	S )CPUBlockStatusblock_idc                 V    t                                                       || _        d S N)super__init__r   )selfr   	__class__s     s/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/vllm/v1/kv_offload/backends/cpu.pyr   zCPUBlockStatus.__init__   s$         )
__name__
__module____qualname__r   _fields_ctypesc_int64intr   __classcell__r   s   @r   r   r      s[        #
FN'C&DDH! ! ! ! ! ! ! ! ! ! !r   r   c                        e Zd Zdedef fdZd Zdee         dee         fdZ	defd	Z
dee         d
ee         defdZ xZS )
CPUBackend
block_size
num_blocksc                     t                                          |t          j                               || _        d| _        g | _        d S )N)r    mediumr   )r   r   r	   r#   r!   num_allocated_blocksallocated_blocks_free_list)r   r    r!   r   s      r   r   zCPUBackend.__init__   sG    J7G7N7P7PQQQ))*!57'''r   c                 J    t          | j                  | j        z   | j        z
  S r   )lenr%   r!   r$   )r   s    r   get_num_free_blockszCPUBackend.get_num_free_blocks   s*    /00o'(	
r   block_hashesreturnc                    t          t          |          | j        | j        z
            }t          |          |z
  }t          | j                  |k    sJ g }t          |          D ]9}|                    t          | j                             | xj        dz  c_        :t          |          D ]=}| j                                        }|                    t          |                     >|S )N   )	minr'   r!   r$   r%   rangeappendr   pop)r   r)   num_fresh_blocksnum_reused_blocksblocks_r   s          r   allocate_blockszCPUBackend.allocate_blocks#   s    t1JJ
 
  --0@@42337HHHHH %''(( 	+ 	+AMM.)BCCDDD%%*%%% ()) 	4 	4A6::<<HMM.223333r   blockc                 r    t          |t                    sJ | j                            |j                   d S r   )
isinstancer   r%   r/   r   )r   r6   s     r   freezCPUBackend.free7   s7    %00000'..u~>>>>>r   r3   c                 4    t          d |D                       S )Nc                     g | ]	}|j         
S  )r   ).0r6   s     r   
<listcomp>z2CPUBackend.get_load_store_spec.<locals>.<listcomp>>   s     D D DE D D Dr   r   )r   r)   r3   s      r   get_load_store_speczCPUBackend.get_load_store_spec;   s!       D DV D D DEEEr   )r   r   r   r   r   r(   listr   r   r5   r9   r   r   r?   r   r   s   @r   r   r      s        83 8C 8 8 8 8 8 8
 
 
DO [@Q    (?+ ? ? ? ?F$Y/F9A+9NF	F F F F F F F Fr   r   )r   collections.abcr   vllm.v1.core.kv_cache_utilsr   vllm.v1.kv_offload.abstractr   vllm.v1.kv_offload.backendr   r   vllm.v1.kv_offload.mediumsr	   r   r   r<   r   r   <module>rF      s     $ $ $ $ $ $ 1 1 1 1 1 1 5 5 5 5 5 5 ; ; ; ; ; ; ; ; 7 7 7 7 7 7! ! ! ! ![ ! ! !*F *F *F *F *F *F *F *F *F *Fr   