nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/zeros_op.py

@@ -0,0 +1,20 @@
+"""Zeros — pure @triton.jit kernel.
+
+Fills a tensor with zeros. Specialization of fill.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def zeros_kernel(
+    output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Fill output with zeros."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id: the product itself wraps past 2^31 elements (register 58)
+    mask = offset < n_elements
+    zero = tl.zeros([BLOCK_SIZE], dtype=output_ptr.type.element_ty)
+    tl.store(output_ptr + offset, zero, mask=mask)

nbxtensor/quant/__init__.py

@@ -0,0 +1 @@
+"""Quantized tensors."""

nbxtensor/quant/quantized_tensor.py

@@ -0,0 +1,187 @@
+"""QuantizedTensor — the engine-side handle for an encoded weight.
+
+A weight stored under the int4-g128-asym encoding rides the container
+as a triplet (`X.qweight` int32 [K//8, N], `X.scales` / `X.qmins`
+fp16 [K//group, N], K = in_features). At weight-load the triplet is
+assembled into ONE QuantizedTensor registered under the graph's
+expected key (`X.weight`), so binding, placement accounting and
+lifecycle stay untouched — compute wrappers detect the type and route:
+
+  aten::t   -> returns self with the view flag flipped. The packed
+               layout is ALREADY the transposed orientation [K, N]
+               (the dequant-GEMV family's natural layout), so the
+               graph's `t(weight [N, K])` is a zero-cost marker.
+  aten::mm  -> M == 1: fused dequant-GEMV (byte-gated family);
+               M > 1: dequantize to dense fp32 via the family's
+               standalone kernel, then the normal mm path (prefill is
+               compute-bound; the transient dense buffer frees at op
+               end).
+
+Lives in kernels/ beside NBXTensor (wrappers route on this type; triton/ consumes it). R33: NBXTensor members only, zero torch. The logical metadata
+(shape/dtype) mirrors what the dense weight would expose so generic
+binding/accounting code needs no special cases.
+"""
+
+from typing import Tuple
+
+from nbxtensor.tensor import NBXTensor
+from nbxtensor.dtype.nbx_dtype import NBXDtype
+
+ENCODING = "int4-g128-asym"
+GROUP = 128
+# The leaves an encoded `X.weight` is stored under, in fold order
+# (packed values, per-group scales, per-group minimums).
+STORAGE_LEAVES = ("qweight", "scales", "qmins")
+PACK = 8
+
+
+class QuantizedTensor:
+    """int4-g128-asym weight triplet under one graph-facing handle.
+
+    `logical_shape` is the DENSE weight's shape as the graph traced it
+    ([out_features, in_features]); `transposed=True` marks the
+    graph-side aten::t view ([in, out] — the packed layout's natural
+    orientation)."""
+
+    __slots__ = ("qweight", "scales", "qmins", "logical_shape",
+                 "transposed")
+
+    def __init__(self, qweight: NBXTensor, scales: NBXTensor,
+                 qmins: NBXTensor, logical_shape: Tuple[int, int],
+                 transposed: bool = False):
+        self.qweight = qweight
+        self.scales = scales
+        self.qmins = qmins
+        self.logical_shape = tuple(logical_shape)
+        self.transposed = transposed
+
+    # ── graph-facing metadata (mirrors the dense weight) ──
+    @property
+    def shape(self) -> Tuple[int, int]:
+        n, k = self.logical_shape
+        return (k, n) if self.transposed else (n, k)
+
+    @property
+    def ndim(self) -> int:
+        return 2
+
+    @property
+    def dtype(self):
+        return NBXDtype.float16  # the dense weight's traced dtype class
+
+    @property
+    def _dtype(self):
+        return NBXDtype.float16
+
+    @property
+    def _shape(self) -> Tuple[int, int]:
+        return self.shape
+
+    @property
+    def nbx_dtype(self):
+        return NBXDtype.float16
+
+    @property
+    def device(self):
+        return self.qweight.device
+
+    @property
+    def _device(self):
+        return self.qweight._device
+
+    @property
+    def _device_idx(self):
+        return self.qweight._device_idx
+
+    @property
+    def _nbytes(self) -> int:
+        return (self.qweight._nbytes + self.scales._nbytes
+                + self.qmins._nbytes)
+
+    def numel(self) -> int:
+        n, k = self.logical_shape
+        return n * k
+
+    def data_ptr(self) -> int:
+        """Identity pointer for cache fingerprints (the qweight buffer
+        — stable for the triplet's lifetime)."""
+        return self.qweight.data_ptr()
+
+    def pin_host(self) -> "QuantizedTensor":
+        """The three host parts pinned (zero3 offload of an int4 weight)."""
+        return QuantizedTensor(self.qweight.pin_host(), self.scales.pin_host(), self.qmins.pin_host(),
+                               self.logical_shape, self.transposed)
+
+    # ── the movers: a weight is moved by ASKING THE WEIGHT ──
+    #
+    # Every placement path in the engine moves a weight by calling a mover on
+    # the tensor itself (`strategies/base.py`, `strategies/zero3.py`,
+    # `strategies/triton/base.py`): dense weights answer, and an encoded one
+    # that cannot answer breaks the first time a plan streams or offloads.
+    # Measured 2026-09-16 on the vitrine's first LLM request:
+    # Qwen3-Coder-30B-A3B-Instruct-int4g128 died at
+    # `'QuantizedTensor' object has no attribute 'to_cuda_async'` — the plan
+    # was right, the encoding was right, the handle was missing three methods.
+    # The triplet moves as a triplet; the metadata (logical shape, view flag)
+    # is not a property of where the bytes live, so it rides unchanged.
+
+    def to_cuda(self, device_idx: int = 0) -> "QuantizedTensor":
+        """The three parts on the given card, the triplet rebuilt."""
+        return QuantizedTensor(self.qweight.to_cuda(device_idx),
+                               self.scales.to_cuda(device_idx),
+                               self.qmins.to_cuda(device_idx),
+                               self.logical_shape, self.transposed)
+
+    def to_cuda_async(self, device_idx: int = 0, stream: int = 0) -> "QuantizedTensor":
+        """The three parts enqueued on `stream` (the streaming strategies'
+        mover). Sequencing is the caller's, exactly as for a dense weight:
+        the returned triplet is readable once the stream is waited on."""
+        return QuantizedTensor(self.qweight.to_cuda_async(device_idx, stream=stream),
+                               self.scales.to_cuda_async(device_idx, stream=stream),
+                               self.qmins.to_cuda_async(device_idx, stream=stream),
+                               self.logical_shape, self.transposed)
+
+    def to_cpu(self, pinned: bool = False) -> "QuantizedTensor":
+        """The three parts back on the host (offload), pinned if asked."""
+        return QuantizedTensor(self.qweight.to_cpu(pinned=pinned),
+                               self.scales.to_cpu(pinned=pinned),
+                               self.qmins.to_cpu(pinned=pinned),
+                               self.logical_shape, self.transposed)
+
+    @property
+    def _pinned(self) -> bool:
+        return bool(getattr(self.qweight, "_pinned", False))
+
+    def t(self) -> "QuantizedTensor":
+        return QuantizedTensor(self.qweight, self.scales, self.qmins,
+                               self.logical_shape,
+                               transposed=not self.transposed)
+
+    def __repr__(self) -> str:
+        return (f"QuantizedTensor({ENCODING}, logical={self.logical_shape}, "
+                f"transposed={self.transposed}, "
+                f"packed={tuple(self.qweight.shape)})")
+
+
+def assemble_quantized(weights: dict) -> int:
+    """Fold `.qweight`/`.scales`/`.qmins` triplets in a loaded weights
+    dict into QuantizedTensor entries under the graph key `<base>.weight`.
+    Mutates the dict in place; returns the number of assembled weights.
+    ZERO FALLBACK: an incomplete triplet raises."""
+    q_leaf, s_leaf, m_leaf = STORAGE_LEAVES
+    bases = [k[: -len("." + q_leaf)] for k in weights
+             if k.endswith("." + q_leaf)]
+    for base in bases:
+        qw = weights.pop(f"{base}.{q_leaf}")
+        try:
+            sc = weights.pop(f"{base}.{s_leaf}")
+            mn = weights.pop(f"{base}.{m_leaf}")
+        except KeyError as e:
+            raise RuntimeError(
+                f"ZERO FALLBACK: incomplete quantized triplet for "
+                f"'{base}' — missing {e}. The build gate should have "
+                f"refused this artifact; re-build the variant.") from e
+        kp, n = qw.shape
+        weights[base + ".weight"] = QuantizedTensor(
+            qw, sc, mn, logical_shape=(n, kp * PACK))
+    return len(bases)

nbxtensor/random/__init__.py

@@ -0,0 +1 @@
+"""Random: seeded streams (the vendor CPU generator's MT19937) and torch-free draws."""

nbxtensor/random/mt_randn.py

@@ -0,0 +1,154 @@
+"""ONE draw for both engines: the vendor's CPU draw, torch-free.
+
+The owner's rule (2026-10-13 02:05): a model's initial noise is one draw, data-driven (the request's
+seed), identical between the compiled and the Triton branch and to the vendor at equal seed. Diffusers
+draws its initial latents on torch's CPU generator — `torch.Generator("cpu").manual_seed(seed)` then
+`randn` — which is MT19937 seeded by `init_genrand(seed)` (the low 32 bits), uniforms made of the 24
+high mantissa bits of a 32-bit word, and ATen's `normal_fill`: Box-Muller over blocks of 16 (the first
+eight uniforms of a block are the radii, the last eight the angles), the tail of `n % 16` values
+recomputed from 16 FRESH uniforms. This module reproduces that stream without torch.
+
+What is bit-exact: the uniform stream (every word of MT19937, every 24-bit mantissa) and the pairing.
+What is within 1 ulp: the transcendentals. ATen computes `sqrt(-2 log u1)`, `cos`, `sin` through its
+vectorised Sleef kernels (`logf_u10`, `sinf/cosf_u35`, AVX512 here, NEON on Apple), which are not
+correctly rounded and differ per architecture; this module computes them in float64 and rounds once,
+which is correctly rounded everywhere. Measured against torch 2.14 on an AVX512 Xeon over 100 000 values
+per seed (seeds 7, 42, 2^31-1; 2026-10-13): the uniform stream bit-equal, the normals at most 3 ulp
+apart (Sleef's u35 bound is 3.5), 98.8 % within 1 ulp, 61.5 % bit-equal, max |diff| 4.8e-7. The draw is
+therefore identical between our two branches and on every platform, and within 4 ulp of the vendor's.
+
+No torch, no device: numpy only (R33 holds; the compiled branch wraps the array in a tensor).
+"""
+from __future__ import annotations
+
+import math
+from typing import Iterable, Tuple
+
+import numpy as np
+
+_N, _M = 624, 397
+_MATRIX_A = np.uint32(0x9908B0DF)
+_UPPER, _LOWER = np.uint32(0x80000000), np.uint32(0x7FFFFFFF)
+
+
+class MT19937:
+    """The generator as torch's CPU `Generator` runs it: `manual_seed(s)` = `init_genrand(s & 0xffffffff)`."""
+
+    def __init__(self, seed: int) -> None:
+        mt = np.zeros(_N, dtype=np.uint32)
+        mt[0] = np.uint32(int(seed) & 0xFFFFFFFF)
+        for i in range(1, _N):   # the recurrence is sequential by definition; 624 steps once per run
+            prev = int(mt[i - 1])
+            mt[i] = np.uint32((1812433253 * (prev ^ (prev >> 30)) + i) & 0xFFFFFFFF)
+        self._mt = mt
+        self._i = _N
+        self._cached = None   # the second Box-Muller sample of the scalar path, as torch caches it
+
+    def _twist(self) -> None:
+        """The reference twist, vectorised in the order it reads: `mt[k] = mt[(k + M) % N] ^ f(y_k)` with
+        `y_k` over the ORIGINAL `mt[k], mt[k+1]` for every k but the last. The second segment reads entries
+        rewritten earlier in the same pass, so it runs in chunks of N-M, each reading only chunks already
+        produced (a single vector expression read uninitialised values — the stream departed at word 454,
+        2026-10-13)."""
+        mt = self._mt
+        y = (mt & _UPPER) | (np.roll(mt, -1) & _LOWER)
+        shifted = (y >> np.uint32(1)) ^ np.where(y & np.uint32(1), _MATRIX_A, np.uint32(0))
+        new = np.empty_like(mt)
+        s = _N - _M                                   # 227: the stride a segment reads behind itself
+        new[:s] = mt[_M:] ^ shifted[:s]
+        k = s
+        while k < _N - 1:
+            e = min(k + s, _N - 1)
+            new[k:e] = new[k - s : e - s] ^ shifted[k:e]
+            k = e
+        y_last = (mt[_N - 1] & _UPPER) | (new[0] & _LOWER)
+        new[_N - 1] = new[_M - 1] ^ ((y_last >> np.uint32(1)) ^ (_MATRIX_A if int(y_last) & 1 else np.uint32(0)))
+        self._mt = new
+        self._i = 0
+
+    def words(self, n: int) -> np.ndarray:
+        """`n` tempered 32-bit words, in stream order."""
+        out = np.empty(n, dtype=np.uint32)
+        k = 0
+        while k < n:
+            if self._i >= _N:
+                self._twist()
+            take = min(n - k, _N - self._i)
+            y = self._mt[self._i : self._i + take].copy()
+            y ^= y >> np.uint32(11)
+            y ^= (y << np.uint32(7)) & np.uint32(0x9D2C5680)
+            y ^= (y << np.uint32(15)) & np.uint32(0xEFC60000)
+            y ^= y >> np.uint32(18)
+            out[k : k + take] = y
+            self._i += take
+            k += take
+        return out
+
+    def uniform_f32(self, n: int) -> np.ndarray:
+        """ATen `uniform_real_distribution<float>(0, 1)`: the 24 low bits of each word over 2^24."""
+        return ((self.words(n) & np.uint32((1 << 24) - 1)).astype(np.float64) * (1.0 / (1 << 24))).astype(np.float32)
+
+    def randn(self, n: int) -> np.ndarray:
+        """`n` standard normals as torch's CPU `randn` draws them (ATen `normal_fill`, blocks of 16; a tail of
+        `n % 16` recomputed from 16 fresh uniforms), the transcendentals in float64 rounded once."""
+        if n < 16:
+            return self._randn_scalar(n)
+        u = self.uniform_f32(n)
+        out = np.empty(n, dtype=np.float32)
+        m = n - n % 16
+        if m:
+            _fill16(out, 0, u[:m])
+        if n % 16:
+            _fill16(out, n - 16, self.uniform_f32(16))
+        return out
+
+    def _randn_scalar(self, n: int) -> np.ndarray:
+        """ATen's path for fewer than 16 values (`normal_kernel`'s serial branch): `normal_distribution<double>`
+        over 53-bit uniforms (two words per uniform, `random64`), Box-Muller in double with the second sample
+        of each pair cached in the generator and served next — exactly the vendor's small draws."""
+        out = np.empty(n, dtype=np.float32)
+        i = 0
+        while i < n:
+            if self._cached is not None:
+                out[i] = np.float32(self._cached)
+                self._cached = None
+                i += 1
+                continue
+            # ATen normal_distribution<double>: u1 then u2; r = sqrt(-2 log1p(-u2)); theta = 2 pi u1;
+            # the sine sample cached for the next call, the cosine returned.
+            u1 = self._uniform_f64()
+            u2 = self._uniform_f64()
+            radius = math.sqrt(-2.0 * math.log1p(-u2))
+            theta = 2.0 * math.pi * u1
+            self._cached = radius * math.sin(theta)
+            out[i] = np.float32(radius * math.cos(theta))
+            i += 1
+        return out
+
+    def _uniform_f64(self) -> float:
+        """ATen `uniform_real_distribution<double>(0, 1)`: the low 53 bits of a 64-bit word made of two 32-bit
+        words (`random64()` = `make64BitsFrom32Bits(random(), random())`: the first word high, the second low)."""
+        hi, lo = (int(x) for x in self.words(2))
+        v = ((hi << 32) | lo) & ((1 << 53) - 1)
+        return v * (1.0 / (1 << 53))
+
+
+def _fill16(dst: np.ndarray, base: int, u: np.ndarray) -> None:
+    b = u.reshape(-1, 16)
+    u1 = (np.float32(1.0) - b[:, :8]).astype(np.float64)
+    u2 = b[:, 8:].astype(np.float32)
+    radius = np.sqrt(-2.0 * np.log(u1))
+    theta = (np.float32(2.0) * np.float32(math.pi) * u2).astype(np.float64)   # ATen: 2.0f * pi<double> * u2 → float
+    view = dst[base : base + b.size].reshape(-1, 16)
+    view[:, :8] = (radius * np.cos(theta)).astype(np.float32)
+    view[:, 8:] = (radius * np.sin(theta)).astype(np.float32)
+
+
+def mt_randn(seed: int, shape: Iterable[int]) -> np.ndarray:
+    """A fresh generator seeded with `seed`, its first `prod(shape)` normals reshaped — the vendor's
+    `randn_tensor(shape, generator=torch.Generator("cpu").manual_seed(seed))`, torch-free."""
+    shape: Tuple[int, ...] = tuple(int(s) for s in shape)
+    n = 1
+    for s in shape:
+        n *= s
+    return MT19937(seed).randn(n).reshape(shape)

nbxtensor/random/rng_pin.py

@@ -0,0 +1,59 @@
+"""Gated cross-engine deterministic RNG for noise pinning.
+
+`NBX_FORCE_RAND_SEED=<int>` makes every `rand`/`randn`/`*_like` op draw from a
+single `numpy.RandomState`, so the PyTorch and Triton engines produce IDENTICAL
+noise when they execute the same DAG (same op order + shapes). Two uses:
+
+1. Diagnostic — pin the noise so an op-by-op cross-engine diff is meaningful
+   past the first random op (otherwise every downstream op "diverges" purely
+   because each engine drew different noise, masking real kernel bugs).
+2. Shared-seeded determinism — the production way to make a stochastic vocoder
+   (e.g. chatterbox s3gen CFM + NSF/SineGen source) reproducible and equal
+   across all four execution modes, same discipline as the openaudio dual_ar /
+   VibeVoice diffusion seeded noise.
+
+Zero torch: this module returns numpy arrays. The Triton side wraps them as
+`NBXTensor.from_numpy`; the PyTorch side wraps them as `torch.from_numpy`. The
+draw order/shape must match across engines for the arrays to be bit-identical;
+both run the same component DAG, so they do.
+"""
+
+import os
+import numpy as np
+
+_STATE = None
+_SEED_USED = "<unset>"
+
+
+def pinned_seed():
+    """Return the configured seed (int) or None when pinning is disabled."""
+    v = os.environ.get("NBX_FORCE_RAND_SEED")
+    if v is None or v == "":
+        return None
+    try:
+        return int(v)
+    except ValueError:
+        return 1234
+
+
+def _state():
+    global _STATE, _SEED_USED
+    s = pinned_seed()
+    if _STATE is None or _SEED_USED != s:
+        _STATE = np.random.RandomState(s if s is not None else 1234)
+        _SEED_USED = s
+    return _STATE
+
+
+def _shape(shape):
+    return tuple(int(d) for d in shape)
+
+
+def pinned_normal(shape):
+    """Standard-normal N(0,1) float32 array of the given shape."""
+    return _state().standard_normal(size=_shape(shape)).astype(np.float32)
+
+
+def pinned_uniform(shape):
+    """Uniform [0,1) float32 array of the given shape."""
+    return _state().random_sample(size=_shape(shape)).astype(np.float32)

nbxtensor/random/rng_stream.py

@@ -0,0 +1,118 @@
+"""Run-scoped seeded RNG stream for the Triton random wrappers.
+
+The random wrappers (`rand/randn/normal/*_like` in `kernels/wrappers.py`)
+historically drew each call's Philox seed from Python's `random.randint`,
+which `--seed` never seeds: Triton-mode stochastic sampling (EulerAncestral
+noise, DDIM eta>0 variance) was NON-reproducible by construction, and no
+run-scoped stream contract existed at all.
+
+This module is the Triton-engine mirror of the compiled engine's run-scoped
+`torch.Generator` (see `VariableResolver.sampling_generator`): ONE seeded
+stream drives every stochastic draw of a run, in consumption order. The flow
+handler arms it at execution start from `defaults.seed` (the same data-driven
+source both engines read); each wrapper call then derives an independent
+per-draw kernel seed as splitmix64(run_seed, draw_counter). Draws are
+reproducible per (seed, draw order) and statistically independent across the
+counter. Cross-ENGINE bit-equality is NOT the contract here (the two engines
+use different RNG algorithms); when a bit-identical cross-engine noise diff
+is needed, the `rng_pin` module (NBX_FORCE_RAND_SEED) remains the tool and
+takes precedence in the wrappers.
+
+Zero torch, zero device state: pure Python integers.
+
+THE ONE DRAW (owner, 2026-10-13 02:05): a model's initial noise is ONE draw, data-driven (the run's
+seed), identical between the compiled and the Triton branch and to the vendor at equal seed. Both
+engines take it from `init_draw(shape)` below — the vendor's CPU draw reproduced torch-free
+(`nbxtensor/random/mt_randn.py`: MT19937 + ATen's normal_fill, within 4 ulp of `torch.randn(generator=cpu)`)
+— as the first draw of the run's seeded stream, then the next for every further `init: randn`
+variable, in consumption order. The per-draw Philox seeds above still serve the Triton random
+kernels of the hot loop (scheduler noise); the compiled branch's scheduler draws keep its generator.
+Cross-engine bit-equality of the INITIAL draw is the contract from here on.
+
+EVERY DRAW OF A SEEDED RUN (2026-10-14): the vendor threads ONE generator through every draw of a run — an
+image encoder's posterior sample (`retrieve_latents(vae.encode(x), generator)`), the initial noise, each
+ancestral step's noise (`EulerAncestralDiscreteScheduler.step(..., generator)`). Measured on Allegro-TI2V at
+its gate request: the two engines rendered the same clip at one step (mean |a - b| under 1 of 255 per frame)
+and different content at ten, the ancestral noise being the only difference (torch's device generator in the
+compiled branch, Philox in the Triton branch); on CogVideoX-5b-I2V the first op that departs between the
+kernel oracle and the Triton branch is the posterior sample `vae_encoder/aten.randn_like::0`. So every draw of a
+seeded run — graph random ops and scheduler noise, in both engines — is the run stream's next normal
+(`next_normal`) or uniform (`next_uniform`) values, in consumption order. The per-draw Philox seeds
+(`next_seed`) serve only what is not a seeded run's draw.
+"""
+
+_RUN_SEED = None
+_COUNTER = 0
+_MT = None            # the run's MT19937, armed with the seed; the init draws come from it in order
+
+
+def set_run_seed(seed):
+    """Arm (or disarm with None) the run stream; resets the draw counter and the init-draw generator."""
+    global _RUN_SEED, _COUNTER, _MT
+    _RUN_SEED = int(seed) if seed is not None else None
+    _COUNTER = 0
+    _MT = None
+
+
+def _mt():
+    global _MT
+    if _RUN_SEED is None:
+        raise RuntimeError("rng_stream: a draw while inactive: the run has no seed")
+    if _MT is None:
+        from nbxtensor.random.mt_randn import MT19937
+        _MT = MT19937(_RUN_SEED)
+    return _MT
+
+
+def _count(shape):
+    n = 1
+    for x in shape:
+        n *= int(x)
+    return n
+
+
+def next_normal(shape):
+    """The run stream's next standard normals as a float32 numpy array of `shape` (torch's CPU `randn` for the
+    run's seed, continued): the draw of every seeded random op of either engine."""
+    shape = tuple(int(x) for x in shape)
+    return _mt().randn(_count(shape)).reshape(shape)
+
+
+def next_uniform(shape):
+    """The run stream's next [0, 1) uniforms as a float32 numpy array of `shape` (torch's CPU `rand`, continued)."""
+    shape = tuple(int(x) for x in shape)
+    return _mt().uniform_f32(_count(shape)).reshape(shape)
+
+
+def init_draw(shape):
+    """The run's next initial-noise draw as a float32 numpy array of `shape` — the vendor's CPU draw for
+    the run's seed (`nbxtensor/random/mt_randn.MT19937`), the same bytes in both engines. Refused when no seed is
+    armed: a seedless run draws unseeded in each engine (the historical behaviour), and that is said."""
+    global _MT
+    if _RUN_SEED is None:
+        raise RuntimeError("rng_stream.init_draw() called while inactive: the run has no seed")
+    if _MT is None:
+        from nbxtensor.random.mt_randn import MT19937
+        _MT = MT19937(_RUN_SEED)
+    shape = tuple(int(x) for x in shape)
+    n = 1
+    for x in shape:
+        n *= x
+    return _MT.randn(n).reshape(shape)
+
+
+def active() -> bool:
+    return _RUN_SEED is not None
+
+
+def next_seed() -> int:
+    """Next per-draw kernel seed: splitmix64 of (run_seed, counter), 31-bit."""
+    global _COUNTER
+    if _RUN_SEED is None:
+        raise RuntimeError("rng_stream.next_seed() called while inactive")
+    x = (_RUN_SEED * 0x9E3779B97F4A7C15 + _COUNTER * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF
+    _COUNTER += 1
+    x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF
+    x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF
+    x = x ^ (x >> 31)
+    return int(x & 0x7FFFFFFF)