mirror of
https://github.com/vladmandic/automatic
synced 2026-08-31 17:41:06 +02:00
typing
This commit is contained in:
@@ -8,11 +8,11 @@ core = ov.Core()
|
||||
|
||||
NPU_MUL = 32 # NPU uses FP16 x INT8 -> FP16 instead of INT8 x INT8 -> INT32 and FP16 output overflows
|
||||
OV_DEVICE: str = os.environ.get("SDNQ_OPENVINO_DEVICE", "CPU")
|
||||
OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, tuple[int,int] | None], list[ov.InferRequest, str]] = {}
|
||||
OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, tuple[int,int] | None], tuple[ov.InferRequest, str]] = {}
|
||||
core.set_property(OV_DEVICE, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
|
||||
|
||||
|
||||
def ov_int_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_name: str) -> torch.Tensor:
|
||||
def ov_int_mm(A: torch.CharTensor, B: torch.CharTensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor:
|
||||
C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32)
|
||||
infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True))
|
||||
infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True))
|
||||
|
||||
@@ -75,7 +75,7 @@ def apply_svdquant(weight: torch.FloatTensor, rank: int = 32, niter: int = 8, dt
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def build_hadamard_n2(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None):
|
||||
def build_hadamard_n2(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
|
||||
current_size = 2
|
||||
H = H_N2 = torch.tensor([[1, 1], [1, -1]], dtype=dtype, device=device)
|
||||
while current_size < n:
|
||||
@@ -87,7 +87,7 @@ def build_hadamard_n2(n: int, dtype: torch.dtype | None = None, device: torch.de
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def build_hadamard_n4(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None):
|
||||
def build_hadamard_n4(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
|
||||
current_size = 4
|
||||
H = H_N4 = torch.tensor([[ 1, 1, 1, -1], [ 1, 1, -1, 1], [ 1, -1, 1, 1], [-1, 1, 1, 1]], dtype=dtype, device=device)
|
||||
while current_size < n:
|
||||
@@ -99,7 +99,7 @@ def build_hadamard_n4(n: int, dtype: torch.dtype | None = None, device: torch.de
|
||||
|
||||
|
||||
@devices.inference_context()
|
||||
def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None):
|
||||
def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
|
||||
if math.log(n, 4).is_integer():
|
||||
return build_hadamard_n4(n, device=device, dtype=dtype)
|
||||
elif math.log(n, 2).is_integer():
|
||||
@@ -111,9 +111,9 @@ def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.devic
|
||||
# 256x256 Hadamard matrix is just 256 KB at FP32
|
||||
# And is the exact same matrix on all model layers
|
||||
# So we can safely cache a single one
|
||||
HADAMARD_MATRIX_CACHE = {}
|
||||
HADAMARD_MATRIX_CACHE: dict[tuple[int, torch.device, torch.dtype], torch.FloatTensor] = {}
|
||||
@devices.inference_context()
|
||||
def get_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None):
|
||||
def get_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor:
|
||||
device = devices.normalize_device(device)
|
||||
H_key = (n, device, dtype)
|
||||
H = HADAMARD_MATRIX_CACHE.get(H_key, None)
|
||||
|
||||
Reference in New Issue
Block a user