mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
19 lines
699 B
Python
19 lines
699 B
Python
"""
|
|
DeepInfer cuda utils
|
|
"""
|
|
from __future__ import annotations
|
|
import typing
|
|
__all__: list[str] = ['alloc_tensor', 'export_ipc_tensor', 'import_ipc_tensor']
|
|
def alloc_tensor(size: typing.SupportsInt | typing.SupportsIndex) -> torch.Tensor | None:
|
|
"""
|
|
Allocate a 1D byte tensor by calling cuMemAlloc directly, bypassing torch's caching allocator.
|
|
"""
|
|
def export_ipc_tensor(ptr: typing.SupportsInt | typing.SupportsIndex, nbytes: typing.SupportsInt | typing.SupportsIndex) -> bytes:
|
|
"""
|
|
Export tensor allocated by `alloc_tensor` via CUDA IPC memory.
|
|
"""
|
|
def import_ipc_tensor(data: bytes) -> torch.Tensor | None:
|
|
"""
|
|
Import CUDA IPC memory handle as a tensor.
|
|
"""
|