smaulbrain-core 3.0.0__tar.gz

smaulbrain_core-3.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 NativeMc
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

smaulbrain_core-3.0.0/PKG-INFO

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: smaulbrain-core
+Version: 3.0.0
+Summary: Byte-level recurrent language model with mixture-of-experts, ternary/FP8 quantization, and a factored optimizer, on PyTorch.
+Author: Yash and contributors to SmaulBRAIN
+License-Expression: LicenseRef-SmaulBRAIN
+Project-URL: Homepage, https://github.com/Native-Neo/SmaulBRAIN
+Project-URL: Repository, https://github.com/Native-Neo/SmaulBRAIN
+Project-URL: Issues, https://github.com/Native-Neo/SmaulBRAIN/issues
+Keywords: language-model,mixture-of-experts,moe,quantization,ternary,pytorch,byte-level
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.6
+Requires-Dist: numpy>=1.24
+Provides-Extra: test
+Requires-Dist: pytest>=8; extra == "test"
+Dynamic: license-file

smaulbrain_core-3.0.0/pyproject.toml

@@ -0,0 +1,59 @@
+[build-system]
+requires = ["setuptools>=77"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "smaulbrain-core"
+dynamic = ["version"]
+description = "Byte-level recurrent language model with mixture-of-experts, ternary/FP8 quantization, and a factored optimizer, on PyTorch."
+readme = "README.md"
+requires-python = ">=3.9"
+license = "LicenseRef-SmaulBRAIN"
+license-files = ["LICENSE"]
+authors = [{ name = "Yash and contributors to SmaulBRAIN" }]
+keywords = [
+    "language-model",
+    "mixture-of-experts",
+    "moe",
+    "quantization",
+    "ternary",
+    "pytorch",
+    "byte-level",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "torch>=2.6",
+    "numpy>=1.24",
+]
+
+[project.optional-dependencies]
+test = ["pytest>=8"]
+
+[project.urls]
+Homepage = "https://github.com/Native-Neo/SmaulBRAIN"
+Repository = "https://github.com/Native-Neo/SmaulBRAIN"
+Issues = "https://github.com/Native-Neo/SmaulBRAIN/issues"
+
+[tool.setuptools]
+packages = ["smaulbrain", "smaulbrain.kernels"]
+py-modules = []
+include-package-data = false
+
+[tool.setuptools.dynamic]
+version = { attr = "smaulbrain.__version__" }
+
+[tool.setuptools.package-data]
+smaulbrain = ["*.cpp"]
+"smaulbrain.kernels" = ["*.cpp"]

smaulbrain_core-3.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

smaulbrain_core-3.0.0/smaulbrain/__init__.py

@@ -0,0 +1,3 @@
+"""SmaulBRAIN V2 package."""
+__version__ = "3.0.0"
+CHECKPOINT_VERSION = "smaulbrain-pytorch-v1:1"

smaulbrain_core-3.0.0/smaulbrain/_native_opt.py

@@ -0,0 +1,13 @@
+"""Optional CPU C++ fused SmaulOPT update reached through Torch tensor views."""
+import torch
+from . import _native_opt_numpy as _backend
+
+
+def factored_step(opt,w,g,st,collect_norms):
+    if w.device.type != "cpu" or not w.is_contiguous() or w.ndim != 2 or w.numel() < 1024:
+        return None
+    if any(not isinstance(st[k],torch.Tensor) or st[k].device.type != "cpu" or not st[k].is_contiguous() for k in ("m","R","C")):
+        return None
+    if st["m"].shape != w.shape or st["R"].shape != (w.shape[0],) or st["C"].shape != (w.shape[1],): return None
+    return _backend.factored_step(opt,w.detach().numpy(),g.detach().contiguous().numpy(),
+                                  {k:st[k].detach().numpy() for k in ("m","R","C")},collect_norms)

smaulbrain_core-3.0.0/smaulbrain/_native_opt_numpy.py

@@ -0,0 +1,92 @@
+"""Optional native fused SmaulOPT update (see _optimizer_kernels.cpp).
+
+Self-contained: builds its own cached shared object and never raises. If no
+compiler is available, SMAULBRAIN_NATIVE_OPT=0 is set, or a tensor does not
+qualify (dtype/layout), callers fall back to the NumPy reference path.
+"""
+import ctypes
+import hashlib
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+import numpy as np
+
+from . import cpu as _cpu
+
+_LIB = None
+_TRIED = False
+
+
+def _load():
+    global _LIB, _TRIED
+    if _LIB is not None or _TRIED:
+        return _LIB
+    _TRIED = True
+    if os.environ.get("SMAULBRAIN_NATIVE_OPT", "1") == "0":
+        return None
+    try:
+        source = Path(__file__).with_name("_optimizer_kernels.cpp")
+        compiler = shutil.which(os.environ.get("CXX", "g++")) or shutil.which("c++")
+        if compiler is None or not source.is_file():
+            return None
+        has_avx = bool(_cpu.detect().get("has_avx", False))
+        variant = "-mavx" if has_avx else "-msse2"
+        flags = ["-O3", "-std=c++17", "-shared", "-fPIC", "-fno-fast-math",
+                 "-ffp-contract=off", "-pthread", variant]
+        digest = hashlib.sha256(source.read_bytes() + compiler.encode()
+                                + " ".join(flags).encode()).hexdigest()[:20]
+        cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "smaulbrain"
+        cache.mkdir(parents=True, exist_ok=True)
+        target = cache / f"opt-{variant[1:]}-{digest}.so"
+        if not target.is_file():
+            tmp = cache / f"opt-{variant[1:]}-{digest}-{os.getpid()}.tmp.so"
+            res = subprocess.run([compiler, *flags, str(source), "-o", str(tmp)],
+                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                                 text=True, timeout=90)
+            if res.returncode != 0:
+                return None
+            os.replace(tmp, target)
+        lib = ctypes.CDLL(str(target))
+        fp = ctypes.POINTER(ctypes.c_float)
+        f = ctypes.c_float
+        lib.smaulbrain_smaulopt_factored.argtypes = [
+            fp, fp, fp, fp, fp, ctypes.c_size_t, ctypes.c_size_t,
+            f, f, f, f, f, f, f, ctypes.c_int, ctypes.POINTER(ctypes.c_double)]
+        lib.smaulbrain_smaulopt_factored.restype = ctypes.c_int
+        _LIB = lib
+    except Exception:
+        _LIB = None
+    return _LIB
+
+
+def _ok(a):
+    return (isinstance(a, np.ndarray) and a.dtype == np.float32
+            and a.flags.c_contiguous and a.flags.writeable)
+
+
+def factored_step(opt, w, g, st, collect_norms):
+    """Run the fused update in place. Returns the update RMS (or 0.0) on
+    success, or None if the caller must use the NumPy path."""
+    lib = _load()
+    if lib is None or w.ndim != 2 or w.size < 1024:
+        return None
+    m, R, C = st["m"], st["R"], st["C"]
+    if not (_ok(w) and _ok(m) and _ok(R) and _ok(C)
+            and m.shape == w.shape and R.shape == (w.shape[0],) and C.shape == (w.shape[1],)):
+        return None
+    gf = np.ascontiguousarray(g, dtype=np.float32)
+    if gf.shape != w.shape:
+        return None
+    f32 = np.float32
+    p = lambda a: a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+    sq = ctypes.c_double(0.0)
+    rc = lib.smaulbrain_smaulopt_factored(
+        p(w), p(m), p(gf), p(R), p(C), w.shape[0], w.shape[1],
+        f32(opt.beta1), f32(1.0 - opt.beta1), f32(opt.beta2), f32(1.0 - opt.beta2),
+        f32(opt.eps), f32(opt.lr), f32(opt.lr * opt.weight_decay) if opt.weight_decay else f32(0.0),
+        1 if collect_norms else 0, ctypes.byref(sq))
+    if rc != 0:
+        return None
+    return float(np.sqrt(sq.value / w.size)) if collect_norms else 0.0

smaulbrain_core-3.0.0/smaulbrain/_optimizer_kernels.cpp

@@ -0,0 +1,265 @@
+// Fused SmaulOPT update for factored (2-D) parameters.
+//
+// Mirrors smaulopt.SmaulOPT.step() for a matrix W (rows x cols):
+//   m <- b1*m + (1-b1)*g
+//   R <- b2*R + (1-b2)*mean_j(g^2),  C <- b2*C + (1-b2)*mean_i(g^2)
+//   upd = m / (sqrt(max(R_i*C_j/mean(R), 0)) + eps)
+//   W  <- W - (lr*wd)*W ;  W <- W - lr*upd
+// in two streaming passes over the matrix (the NumPy path makes ~10).
+// All per-element arithmetic is IEEE float32 in the same operation order as
+// the NumPy path (no FMA, no fast-math); only the row/column means differ,
+// they accumulate in float64. Reductions use fixed 16-row blocks combined in
+// block order, so results do not depend on the thread count.
+#include <algorithm>
+#include <atomic>
+#include <cmath>
+#include <condition_variable>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <functional>
+#include <mutex>
+#include <thread>
+#include <vector>
+#if defined(__AVX__)
+#include <immintrin.h>
+#endif
+
+namespace {
+constexpr std::size_t kBlockRows = 16;
+
+// Persistent worker pool (same pattern as _ternary_kernels.cpp). Creating
+// threads per optimizer call costs more than the update itself at matrix
+// sizes this project targets, so workers are started once per process. Worker
+// count = SMAULBRAIN_THREADS if set (>=1), else hardware_concurrency capped at
+// 16; SMAULBRAIN_THREADS=1 disables threading.
+class Pool {
+public:
+    // Intentionally leaked: detached workers must outlive static destruction.
+    static Pool& get() { static Pool* p = new Pool(); return *p; }
+    unsigned size() const noexcept { return nthreads_; }
+
+    void run(std::size_t n_items, const std::function<void(std::size_t)>& fn) {
+        if (n_items == 0) return;
+        if (nthreads_ <= 1 || n_items == 1) {
+            for (std::size_t i = 0; i < n_items; ++i) fn(i);
+            return;
+        }
+        std::lock_guard<std::mutex> run_lock(run_mu_);
+        fn_ = &fn;
+        n_items_ = n_items;
+        next_.store(0, std::memory_order_relaxed);
+        const unsigned helpers = static_cast<unsigned>(
+            std::min<std::size_t>(nthreads_ - 1, n_items - 1));
+        remaining_.store(helpers, std::memory_order_relaxed);
+        wanted_.store(helpers, std::memory_order_relaxed);
+        {
+            std::lock_guard<std::mutex> lk(mu_);
+            generation_.fetch_add(1, std::memory_order_release);
+        }
+        cv_.notify_all();
+        work();
+        while (remaining_.load(std::memory_order_acquire) != 0) {
+            for (int k = 0; k < 64; ++k) cpu_relax();
+            if (remaining_.load(std::memory_order_acquire) == 0) break;
+            std::this_thread::yield();
+        }
+    }
+
+private:
+    Pool() {
+        unsigned n = std::thread::hardware_concurrency();
+        if (n == 0) n = 1;
+        if (const char* e = std::getenv("SMAULBRAIN_THREADS")) {
+            const long v = std::strtol(e, nullptr, 10);
+            if (v >= 1) n = static_cast<unsigned>(v);
+        }
+        nthreads_ = std::min(n, 16u);
+        for (unsigned i = 1; i < nthreads_; ++i)
+            workers_.emplace_back([this] { loop(); });
+        for (auto& t : workers_) t.detach();  // process exit must not block on workers
+    }
+    static inline void cpu_relax() noexcept {
+#if defined(__SSE2__)
+        _mm_pause();
+#endif
+    }
+    void work() {
+        for (;;) {
+            const std::size_t i = next_.fetch_add(1, std::memory_order_relaxed);
+            if (i >= n_items_) break;
+            (*fn_)(i);
+        }
+    }
+    void loop() {
+        unsigned seen = 0;
+        for (;;) {
+            // Spin briefly for the next job, then sleep on the condvar.
+            bool got = false;
+            for (int k = 0; k < 20000; ++k) {
+                if (generation_.load(std::memory_order_acquire) != seen) { got = true; break; }
+                cpu_relax();
+            }
+            if (!got) {
+                std::unique_lock<std::mutex> lk(mu_);
+                cv_.wait(lk, [&] { return generation_.load(std::memory_order_acquire) != seen; });
+            }
+            seen = generation_.load(std::memory_order_acquire);
+            // Only `wanted_` helpers participate in each job.
+            unsigned w = wanted_.load(std::memory_order_relaxed);
+            while (w > 0 && !wanted_.compare_exchange_weak(w, w - 1, std::memory_order_acq_rel)) {}
+            if (w == 0) continue;
+            work();
+            remaining_.fetch_sub(1, std::memory_order_acq_rel);
+        }
+    }
+    unsigned nthreads_ = 1;
+    std::vector<std::thread> workers_;
+    std::mutex mu_, run_mu_;
+    std::condition_variable cv_;
+    std::atomic<unsigned> generation_{0}, remaining_{0}, wanted_{0};
+    std::atomic<std::size_t> next_{0};
+    const std::function<void(std::size_t)>* fn_ = nullptr;
+    std::size_t n_items_ = 0;
+};
+
+// Dynamic work distribution over [0, nblocks). Callers must make each block's
+// output independent of which thread runs it, and combine per-block results in
+// block order (as the passes below do), so results stay deterministic.
+template <class F>
+void parallel_blocks(std::size_t nblocks, bool allow_threads, F&& fn) {
+    Pool& pool = Pool::get();
+    if (!allow_threads || nblocks == 0 || pool.size() <= 1) {
+        for (std::size_t b = 0; b < nblocks; ++b) fn(b);
+        return;
+    }
+    std::function<void(std::size_t)> job = [&fn](std::size_t b) { fn(b); };
+    pool.run(nblocks, job);
+}
+}  // namespace
+
+extern "C" int smaulbrain_smaulopt_factored(
+    float* w, float* m, const float* g, float* R, float* C,
+    std::size_t rows, std::size_t cols,
+    float beta1, float omb1, float beta2, float omb2, float eps, float lr,
+    float decay, int want_norm, double* upd_sq_sum) noexcept {
+    if (!w || !m || !g || !R || !C || rows == 0 || cols == 0) return 1;
+    try {
+        const std::size_t nblocks = (rows + kBlockRows - 1) / kBlockRows;
+        const bool threaded = rows * cols >= (std::size_t{1} << 18);
+
+        // ---- pass 1: moment update + g^2 row/column statistics -------------
+        std::vector<double> row_sum(rows, 0.0);
+        std::vector<float> col_part(nblocks * cols, 0.0f);
+        parallel_blocks(nblocks, threaded, [&](std::size_t blk) {
+            const std::size_t i0 = blk * kBlockRows, i1 = std::min(rows, i0 + kBlockRows);
+            float* cp = col_part.data() + blk * cols;
+            for (std::size_t i = i0; i < i1; ++i) {
+                const float* gr = g + i * cols;
+                float* mr = m + i * cols;
+                double rs = 0.0;
+                std::size_t j = 0;
+#if defined(__AVX__)
+                const __m256 b1 = _mm256_set1_ps(beta1), o1 = _mm256_set1_ps(omb1);
+                __m256d acc = _mm256_setzero_pd();
+                for (; j + 8 <= cols; j += 8) {
+                    const __m256 gv = _mm256_loadu_ps(gr + j);
+                    const __m256 mv = _mm256_add_ps(_mm256_mul_ps(_mm256_loadu_ps(mr + j), b1),
+                                                    _mm256_mul_ps(gv, o1));
+                    _mm256_storeu_ps(mr + j, mv);
+                    const __m256 g2 = _mm256_mul_ps(gv, gv);
+                    _mm256_storeu_ps(cp + j, _mm256_add_ps(_mm256_loadu_ps(cp + j), g2));
+                    acc = _mm256_add_pd(acc, _mm256_cvtps_pd(_mm256_castps256_ps128(g2)));
+                    acc = _mm256_add_pd(acc, _mm256_cvtps_pd(_mm256_extractf128_ps(g2, 1)));
+                }
+                alignas(32) double lanes[4];
+                _mm256_store_pd(lanes, acc);
+                rs = (lanes[0] + lanes[1]) + (lanes[2] + lanes[3]);
+#endif
+                for (; j < cols; ++j) {
+                    const float gv = gr[j];
+                    mr[j] = mr[j] * beta1 + gv * omb1;
+                    const float g2 = gv * gv;
+                    cp[j] += g2;
+                    rs += static_cast<double>(g2);
+                }
+                row_sum[i] = rs;
+            }
+        });
+
+        // ---- statistics update (tiny) --------------------------------------
+        std::vector<double> col_sum(cols, 0.0);
+        for (std::size_t blk = 0; blk < nblocks; ++blk) {
+            const float* cp = col_part.data() + blk * cols;
+            for (std::size_t j = 0; j < cols; ++j) col_sum[j] += static_cast<double>(cp[j]);
+        }
+        double r_total = 0.0;
+        for (std::size_t i = 0; i < rows; ++i) {
+            const float row_mean = static_cast<float>(row_sum[i] / static_cast<double>(cols));
+            R[i] = R[i] * beta2 + row_mean * omb2;
+            r_total += static_cast<double>(R[i]);
+        }
+        for (std::size_t j = 0; j < cols; ++j) {
+            const float col_mean = static_cast<float>(col_sum[j] / static_cast<double>(rows));
+            C[j] = C[j] * beta2 + col_mean * omb2;
+        }
+        const double mean_r = r_total / static_cast<double>(rows);
+        const float denom = static_cast<float>(mean_r > 1e-30 ? mean_r : 1e-30);
+
+        // ---- pass 2: update + weight decay (+ optional update norm) --------
+        std::vector<double> blk_sq(nblocks, 0.0);
+        parallel_blocks(nblocks, threaded, [&](std::size_t blk) {
+            const std::size_t i0 = blk * kBlockRows, i1 = std::min(rows, i0 + kBlockRows);
+            double sq = 0.0;
+            for (std::size_t i = i0; i < i1; ++i) {
+                float* wr = w + i * cols;
+                const float* mr = m + i * cols;
+                const float ri = R[i];
+                std::size_t j = 0;
+#if defined(__AVX__)
+                const __m256 rv = _mm256_set1_ps(ri), dv = _mm256_set1_ps(denom);
+                const __m256 ev = _mm256_set1_ps(eps), lv = _mm256_set1_ps(lr);
+                const __m256 cv = _mm256_set1_ps(decay), zero = _mm256_setzero_ps();
+                __m256d acc = _mm256_setzero_pd();
+                for (; j + 8 <= cols; j += 8) {
+                    __m256 v = _mm256_mul_ps(rv, _mm256_loadu_ps(C + j));
+                    v = _mm256_div_ps(v, dv);
+                    v = _mm256_max_ps(zero, v);  // NaN-propagating like np.maximum
+                    v = _mm256_add_ps(_mm256_sqrt_ps(v), ev);
+                    const __m256 u = _mm256_div_ps(_mm256_loadu_ps(mr + j), v);
+                    if (want_norm) {
+                        const __m256 u2 = _mm256_mul_ps(u, u);
+                        acc = _mm256_add_pd(acc, _mm256_cvtps_pd(_mm256_castps256_ps128(u2)));
+                        acc = _mm256_add_pd(acc, _mm256_cvtps_pd(_mm256_extractf128_ps(u2, 1)));
+                    }
+                    __m256 wv = _mm256_loadu_ps(wr + j);
+                    if (decay != 0.0f) wv = _mm256_sub_ps(wv, _mm256_mul_ps(cv, wv));
+                    wv = _mm256_sub_ps(wv, _mm256_mul_ps(u, lv));
+                    _mm256_storeu_ps(wr + j, wv);
+                }
+                alignas(32) double lanes[4];
+                _mm256_store_pd(lanes, acc);
+                sq += (lanes[0] + lanes[1]) + (lanes[2] + lanes[3]);
+#endif
+                for (; j < cols; ++j) {
+                    float v = ri * C[j];
+                    v = v / denom;
+                    v = (v != v) ? v : (v > 0.0f ? v : 0.0f);
+                    v = std::sqrt(v) + eps;
+                    const float u = mr[j] / v;
+                    if (want_norm) sq += static_cast<double>(u * u);
+                    float wv = wr[j];
+                    if (decay != 0.0f) wv = wv - decay * wv;
+                    wr[j] = wv - u * lr;
+                }
+            }
+            blk_sq[blk] = sq;
+        });
+        double total = 0.0;
+        for (double v : blk_sq) total += v;
+        if (upd_sq_sum) *upd_sq_sum = total;
+        return 0;
+    } catch (...) {
+        return 2;
+    }
+}