This commit is contained in:
Disty0
2026-06-19 03:28:09 +03:00
parent bd3eac047f
commit 440aa0d105
2 changed files with 7 additions and 7 deletions
+2 -2
View File
@@ -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))
+5 -5
View File
@@ -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)