o پi@sUddlZddlmZmZddlZiZeeeejfejfe d<dede dejdejfdd Z d d Z ej d d defddZdS)N)DictTuple _cache_bufnamebytesdevicereturncCs8||f}t|}|durtj|tj|d}|t|<|S)N)dtyper)rgettorchemptyuint8)rrrkeybufrD/home/ubuntu/.local/lib/python3.10/site-packages/sgl_kernel/utils.py_get_cache_bufs  rcCst|tjr |dfSd|fS)Nr) isinstancer Tensor)xrrr_to_tensor_scalar_tuple!s r)maxsizecCs"tj}tj|\}}|dkS)N )r cudacurrent_deviceget_device_capability)rmajorminorrrris_arch_support_pdl(s r) functoolstypingrrr rstrrr__annotations__intrr lru_cacheboolrrrrrs