nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/memory/vmm.py

@@ -0,0 +1,118 @@
+"""The device's virtual memory management primitives (CUDA's driver API: cuMemAddressReserve, cuMemCreate, cuMemMap,
+cuMemSetAccess, cuMemUnmap, cuMemRelease) — what an expandable range (`memory.expandable`) builds on: ONE virtual
+address range per device, physical chunks mapped into it as it grows and unmapped where a free block covers them whole
+(R16, the supervisor's GO 2026-11-02 01:30; the method of torch's `expandable_segments`, PR #96995).
+
+Reached through ctypes on the driver library the backend table names (`driver_libs`); a backend without the API (HIP,
+Metal) or a device without the attribute (`CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED`) answers
+`supported() == False` and the pool keeps its driver segments. Every call returns the driver's status code; the caller
+names a failure. Measured on V100 (registres/benchmarks/vmm_v100.md, 2026-11-02): the granularity is 2 MiB; a map of a
+chunk (create + map + set access) costs what an ordinary allocation of the same size costs (~0.36 ms), an unmap ~0.1 ms.
+
+Torch-free, numpy-free.
+"""
+from __future__ import annotations
+
+import ctypes
+from typing import Optional, Tuple
+
+#: cuMemAllocationProp.type: CU_MEM_ALLOCATION_TYPE_PINNED; .location.type: CU_MEM_LOCATION_TYPE_DEVICE
+_ALLOC_PINNED, _LOC_DEVICE = 1, 1
+#: CU_MEM_ACCESS_FLAGS_PROT_READWRITE
+_PROT_RW = 3
+#: CU_MEM_ALLOC_GRANULARITY_MINIMUM
+_GRAN_MINIMUM = 0
+#: CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED
+_ATTR_VMM = 102
+
+
+class _Loc(ctypes.Structure):
+    _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)]
+
+
+class _AllocFlags(ctypes.Structure):
+    _fields_ = [("compressionType", ctypes.c_ubyte), ("gpuDirectRDMACapable", ctypes.c_ubyte),
+                ("usage", ctypes.c_ushort), ("reserved", ctypes.c_ubyte * 4)]
+
+
+class _Prop(ctypes.Structure):
+    _fields_ = [("type", ctypes.c_int), ("requestedHandleTypes", ctypes.c_int), ("location", _Loc),
+                ("win32HandleMetaData", ctypes.c_void_p), ("allocFlags", _AllocFlags)]
+
+
+class _Access(ctypes.Structure):
+    _fields_ = [("location", _Loc), ("flags", ctypes.c_int)]
+
+
+class Vmm:
+    """One device's VMM calls. `lib`: the driver library (ctypes CDLL); `device`: the device ordinal the process sees."""
+
+    def __init__(self, lib, device: int) -> None:
+        self.lib = lib
+        self.device = int(device)
+        self._prop = _Prop()
+        self._prop.type = _ALLOC_PINNED
+        self._prop.location = _Loc(_LOC_DEVICE, self.device)
+        self._access = _Access(_Loc(_LOC_DEVICE, self.device), _PROT_RW)
+
+    @staticmethod
+    def load(libs) -> Optional[object]:
+        """The first driver library of `libs` that loads, initialised; None when none does."""
+        for name in libs or ():
+            try:
+                lib = ctypes.CDLL(name)
+            except OSError:
+                continue
+            if lib.cuInit(0) == 0:
+                return lib
+        return None
+
+    def supported(self) -> bool:
+        """The device answers the VMM attribute with 1."""
+        dev, val = ctypes.c_int(), ctypes.c_int()
+        if self.lib.cuDeviceGet(ctypes.byref(dev), self.device) != 0:
+            return False
+        if self.lib.cuDeviceGetAttribute(ctypes.byref(val), _ATTR_VMM, dev) != 0:
+            return False
+        return val.value == 1
+
+    def granularity(self) -> Tuple[int, int]:
+        """(rc, the minimum allocation granularity in bytes)."""
+        g = ctypes.c_size_t()
+        rc = self.lib.cuMemGetAllocationGranularity(ctypes.byref(g), ctypes.byref(self._prop), _GRAN_MINIMUM)
+        return rc, int(g.value)
+
+    def reserve(self, nbytes: int) -> Tuple[int, int]:
+        """(rc, the base of a reserved virtual range of `nbytes`)."""
+        va = ctypes.c_uint64()
+        rc = self.lib.cuMemAddressReserve(ctypes.byref(va), ctypes.c_size_t(nbytes), ctypes.c_size_t(0),
+                                          ctypes.c_uint64(0), ctypes.c_ulonglong(0))
+        return rc, int(va.value)
+
+    def unreserve(self, va: int, nbytes: int) -> int:
+        return self.lib.cuMemAddressFree(ctypes.c_uint64(va), ctypes.c_size_t(nbytes))
+
+    def create(self, nbytes: int) -> Tuple[int, int]:
+        """(rc, a physical allocation handle of `nbytes` on the device)."""
+        h = ctypes.c_ulonglong()
+        rc = self.lib.cuMemCreate(ctypes.byref(h), ctypes.c_size_t(nbytes), ctypes.byref(self._prop),
+                                  ctypes.c_ulonglong(0))
+        return rc, int(h.value)
+
+    def map(self, va: int, nbytes: int, handle: int) -> int:
+        """Map `handle` at `va` and open it to the device; on a failed access grant the mapping is undone."""
+        rc = self.lib.cuMemMap(ctypes.c_uint64(va), ctypes.c_size_t(nbytes), ctypes.c_size_t(0),
+                               ctypes.c_ulonglong(handle), ctypes.c_ulonglong(0))
+        if rc != 0:
+            return rc
+        rc = self.lib.cuMemSetAccess(ctypes.c_uint64(va), ctypes.c_size_t(nbytes), ctypes.byref(self._access),
+                                     ctypes.c_size_t(1))
+        if rc != 0:
+            self.lib.cuMemUnmap(ctypes.c_uint64(va), ctypes.c_size_t(nbytes))
+        return rc
+
+    def unmap(self, va: int, nbytes: int) -> int:
+        return self.lib.cuMemUnmap(ctypes.c_uint64(va), ctypes.c_size_t(nbytes))
+
+    def release(self, handle: int) -> int:
+        return self.lib.cuMemRelease(ctypes.c_ulonglong(handle))

nbxtensor/ops/__init__.py

@@ -0,0 +1 @@
+"""Ops: the Triton kernels, one file per op (pure @triton.jit compute). No torch in this package."""

nbxtensor/ops/_autotune_policy.py

@@ -0,0 +1,90 @@
+"""NBX_DISABLE_AUTOTUNE — diagnostic env var that forces all autotuned
+matmul-class kernels to use a single fixed config.
+
+Purpose (P-SANA-4KPX-RUNTIME 2026-05-07): isolate autotune as the
+cause of shape-dependent numerical divergence at Sana 4Kpx scale.
+When `NBX_DISABLE_AUTOTUNE=1`, the autotune `configs=[...]` list is
+filtered down to a single `Config` known to be Volta-viable, so the
+autotuner has no degrees of freedom and the same bit-pattern of
+accumulation runs at every shape. If running with this flag yields a
+coherent PNG that the autotune-enabled run does not, autotune is
+confirmed as the cause.
+
+This is **not for production** — autotune is required to reach
+cuBLAS perf parity per CLAUDE.md autotune doctrine. Use only as a
+diagnostic / temporary workaround until a buggy config is removed
+from the search space or its kernel is fixed.
+
+Universal across hardware (R23): the fixed config is `BLOCK_M=64,
+BLOCK_N=64, BLOCK_K=32, num_warps=4, num_stages=2` for matmul-class
+kernels — Volta-viable per CLAUDE.md autotune policy and within
+sm_80+ SMEM headroom too.
+"""
+
+import os
+
+NBX_DISABLE_AUTOTUNE = os.environ.get("NBX_DISABLE_AUTOTUNE", "0") == "1"
+
+# Diagnostic env var (P-SANA-4KPX-RUNTIME 2026-05-08): forces all
+# half-precision (fp16/bf16) inputs to mm/bmm/addmm to be upcast to
+# fp32 BEFORE entering the matmul kernel, regardless of hardware
+# capability or M size. Also forces output dtype to fp32 and
+# IEEE_PRECISION=True. Used to isolate whether the residual Sana 4Kpx
+# bug is caused by a dtype-intermediate path (cuDNN-vs-NBX
+# accumulation precision difference) that the existing wrapper rules
+# don't cover for bf16 inputs on Volta. NOT for production - costs
+# 2x VRAM and disables HMMA. Once a definitive verdict is reached
+# (Cas A: bug fixed -> isolate the specific dtype path that needs
+# correction; Cas B: bug persists -> dtype intermediate eliminated as
+# suspect, focus shifts to reduction order or layout/stride).
+NBX_FORCE_FP32_ACCUM = os.environ.get("NBX_FORCE_FP32_ACCUM", "0") == "1"
+
+
+def maybe_pin_single(configs, predicate):
+    """If `NBX_DISABLE_AUTOTUNE=1`, return a single-element list with
+    the FIRST config matching `predicate`. Otherwise return `configs`
+    unchanged.
+
+    `predicate` is a callable `Config -> bool`. The first match in the
+    `configs` list is selected, so callers should prefer ordering with
+    the desired Volta-safe target near the top (or use a unique
+    predicate).
+    """
+    if not NBX_DISABLE_AUTOTUNE:
+        return configs
+    for c in configs:
+        if predicate(c):
+            return [c]
+    # Fallback: keep a single arbitrary config if the predicate finds
+    # nothing (e.g. depthwise has different field names). Caller must
+    # adjust predicate per kernel.
+    return configs[:1]
+
+
+def is_matmul_pinned(c):
+    """Match the matmul/bmm/addmm Volta-safe target:
+    BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, num_warps=4, num_stages=2.
+    """
+    kw = c.kwargs
+    return (kw.get('BLOCK_M') == 64 and kw.get('BLOCK_N') == 64
+            and kw.get('BLOCK_K') == 32
+            and c.num_warps == 4 and c.num_stages == 2)
+
+
+def is_conv2d_pinned(c):
+    """Match conv2d Volta-safe target:
+    BLOCK_SIZE_BHW=64, BLOCK_SIZE_OUTF=64, BLOCK_SIZE_INF=32, warps=4, stages=2.
+    """
+    kw = c.kwargs
+    return (kw.get('BLOCK_SIZE_BHW') == 64
+            and kw.get('BLOCK_SIZE_OUTF') == 64
+            and kw.get('BLOCK_SIZE_INF') == 32
+            and c.num_warps == 4 and c.num_stages == 2)
+
+
+def is_depthwise_conv2d_pinned(c):
+    """Match depthwise conv2d's pinned config — its fields are the 2-D tile's (BC channels x TH x TW): 2 x 8 x 64,
+    the best measured on SANA-Video's 11200- and 13440-channel 3x3 (V100-SXM2, 2026-11-03).
+    """
+    kw = c.kwargs
+    return kw.get('BC') == 2 and kw.get('TH') == 8 and kw.get('TW') == 64

nbxtensor/ops/_common.py

@@ -0,0 +1,32 @@
+"""Shared activation element functions — pure @triton.jit, no pointers.
+
+These operate on already-loaded values and compose naturally.
+Ported from attorch (MIT license).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def sigmoid(x):
+    """Sigmoid: 1 / (1 + exp(-x))"""
+    return 1.0 / (1.0 + tl.exp(-x))
+
+
+@triton.jit
+def tanh_fn(x):
+    """Tanh via sigmoid: 2 * sigmoid(2x) - 1"""
+    return 2.0 * sigmoid(2.0 * x) - 1.0
+
+
+@triton.jit
+def relu_fn(x):
+    """ReLU: max(0, x)"""
+    return tl.maximum(0, x)
+
+
+@triton.jit
+def relu6_fn(x):
+    """ReLU6: min(max(0, x), 6)"""
+    return tl.minimum(relu_fn(x), 6)