bitsandbytes-cpu-fork 0.50.2.dev4__py3-none-manylinux_2_39_x86_64.whl

bitsandbytes/__init__.py

@@ -0,0 +1,81 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import importlib
+import sys
+
+import torch
+
+from . import _ops, utils
+from .autograd._functions import (
+    MatmulLtState,
+    matmul,
+    matmul_4bit,
+)
+from .backends.cpu import ops as cpu_ops
+from .backends.default import ops as default_ops
+from .nn import modules
+from .optim import adam
+
+# This is a signal for integrations with transformers/diffusers.
+# Eventually we may remove this but it is currently required for compatibility.
+features = {"multi_backend"}
+supported_torch_devices = {
+    "cpu",
+    "cuda",  # NVIDIA/AMD GPU
+    "xpu",  # Intel GPU
+    "hpu",  # Intel Gaudi
+    "npu",  # Ascend NPU
+    "mps",  # Apple Silicon
+}
+
+if torch.cuda.is_available():
+    from .backends.cuda import ops as cuda_ops
+
+if hasattr(torch, "xpu") and torch.xpu.is_available():
+    from .backends.xpu import ops as xpu_ops
+
+if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    from .backends.mps import ops as mps_ops
+
+if importlib.util.find_spec("habana_frameworks") and importlib.util.find_spec("habana_frameworks.torch"):
+    # In case not automatically imported
+    import habana_frameworks.torch
+
+    if hasattr(torch, "hpu") and torch.hpu.is_available():
+        from .backends.hpu import ops as hpu_ops
+
+
+def _import_backends():
+    """
+    Discover and autoload all available backends installed as separate packages.
+    Packages with an entrypoint for "bitsandbytes.backends" will be loaded.
+    Inspired by PyTorch implementation: https://pytorch.org/tutorials/prototype/python_extension_autoload.html
+    """
+    from importlib.metadata import entry_points
+
+    extensions = entry_points(group="bitsandbytes.backends")
+
+    for ext in extensions:
+        try:
+            entry = ext.load()
+            entry()
+        except Exception as e:
+            raise RuntimeError(f"bitsandbytes: failed to load backend {ext.name}: {e}") from e
+
+
+_import_backends()
+
+__pdoc__ = {
+    "libbitsandbytes": False,
+    "optim.optimizer.Optimizer8bit": False,
+    "optim.optimizer.MockArgs": False,
+}
+
+__version__ = "0.50.2.dev4"
+
+# 挂载 gdn_cpu（GDN 融合内核 / patch_transformers 入口），让 bnb.gdn_cpu 可直接访问
+from . import gdn_cpu  # noqa: E402,F401

bitsandbytes/__main__.py

@@ -0,0 +1,4 @@
+if __name__ == "__main__":
+    from bitsandbytes.diagnostics.main import main
+
+    main()

bitsandbytes/_cpu_caps.py

@@ -0,0 +1,110 @@
+"""CPU feature queries used to choose between the AVX2 and AVX-512 CPU libraries.
+
+Why a separate module: two places need the same answer -- `cextension` (the main native
+library loader) and `gdn_cpu` (the fused Gated DeltaNet wrapper). Importing one from the
+other would drag the whole native library in at import time, so the query lives here and
+imports nothing but the standard library; torch is imported lazily, only when it is already
+present in the process.
+
+The AVX-512 backend is a separate artifact because it is NOT safe on AVX2-only CPUs
+(and it is compiled with AVX512-BF16, so it is also unsafe on AVX-512 CPUs without BF16):
+building it with `-mavx512f` is a translation-unit-wide switch, and GCC auto-vectorises
+ordinary kernels (dequantise, optimizer, rmsnorm) with zmm instructions. The library must
+therefore only be loaded after the CPU and OS are known to support AVX-512. On Linux the
+`/proc/cpuinfo` flags are authoritative; on Windows the query goes through torch, which
+already performs the full CPUID + XGETBV check (and exists in any process that can use this
+package). If neither source can answer, the answer is "no" -- refusing to load the AVX-512
+library is always safe, guessing is not.
+"""
+from __future__ import annotations
+
+import os
+
+# avx512_bf16 is required too: the AVX-512 artifact is compiled with
+# -mavx512bf16, so it contains bf16 instructions (auto-vectorised) even outside
+# the explicitly guarded bf16 kernels. CPUs with AVX-512 but no BF16
+# (Skylake-X, Cascade Lake) would SIGILL on them.
+_AVX512_FLAGS = frozenset({"avx512f", "avx512bw", "avx512dq", "avx512vl", "avx512_bf16"})
+
+
+def _env_flag(name: str) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return False
+    return value.strip().lower() not in ("", "0", "false", "no", "off")
+
+
+def _torch_capability() -> str | None:
+    """What torch already knows about this CPU ('AVX2', 'AVX512', 'DEFAULT', ...)."""
+    try:
+        import torch
+    except Exception:  # pragma: no cover - torch is a hard dependency in practice
+        return None
+    backend = getattr(torch, "backends", None)
+    cpu = getattr(backend, "cpu", None)
+    query = getattr(cpu, "get_cpu_capability", None)
+    if query is None:
+        return None
+    try:
+        return str(query()).upper()
+    except Exception:  # pragma: no cover - defensive: a broken torch must not crash import
+        return None
+
+
+def _read_cpuinfo_flags(path: str = "/proc/cpuinfo") -> set[str] | None:
+    """First ``flags`` line of /proc-style cpuinfo, or None when unreadable (non-Linux)."""
+    try:
+        with open(path, "r", encoding="utf-8", errors="replace") as handle:
+            for line in handle:
+                if line.startswith("flags"):
+                    return set(line.split(":", 1)[1].split())
+    except OSError:
+        return None
+    return None
+
+
+def _proc_cpuinfo_supports_avx2(path: str = "/proc/cpuinfo") -> bool:
+    flags = _read_cpuinfo_flags(path)
+    return flags is not None and {"avx2", "fma", "f16c"} <= flags
+
+
+def _proc_cpuinfo_supports_avx512(path: str = "/proc/cpuinfo") -> bool:
+    flags = _read_cpuinfo_flags(path)
+    return flags is not None and _AVX512_FLAGS <= flags
+
+
+def avx512_supported() -> bool:
+    """True only when the CPU + OS can execute the AVX-512 libraries' instructions.
+
+    BNB_CPU_DISABLE_AVX512=1 forces the AVX2 library even on capable hardware (useful when
+    comparing the two backends); BNB_CPU_FORCE_AVX512=1 forces the AVX-512 path (only safe on
+    a machine that really supports it -- it exists so the selection logic can be tested).
+    """
+    if _env_flag("BNB_CPU_DISABLE_AVX512"):
+        return False
+    if _env_flag("BNB_CPU_FORCE_AVX512"):
+        return True
+    flags = _read_cpuinfo_flags()
+    if flags is not None:
+        # /proc/cpuinfo knows about BF16, so it is the authoritative answer on Linux.
+        return _AVX512_FLAGS <= flags
+    # torch reports only the ISA level ("AVX512"), not BF16. Since the shipped AVX-512
+    # library contains bf16 instructions, an ISA-level "yes" alone is not safe there.
+    capability = _torch_capability()
+    if capability is not None:
+        return False
+    return False
+
+
+def avx2_supported() -> bool:
+    """True when the CPU can execute the regular (AVX2+FMA) artifact.
+
+    BNB_CPU_FORCE_GENERIC=1 forces the portable artifact (it runs on any x86-64, so this is
+    safe and makes the three-way selection testable on an AVX2 machine).
+    """
+    if _env_flag("BNB_CPU_FORCE_GENERIC"):
+        return False
+    capability = _torch_capability()
+    if capability is not None:
+        return capability in ("AVX2", "AVX512")
+    return _proc_cpuinfo_supports_avx2()