nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/launch_observer.py

@@ -0,0 +1,54 @@
+"""An optional observer of kernel launches, and a dry-run mode behind it.
+
+An application may install ONE observer (`set_launch_observer`). Every autotuned launch hands its key to
+`record(tuned, key)`, live or dry (an application's record of the keys a run forms). While `active()` answers True the
+library runs dry: an autotuned kernel launches nothing — it calls `shadow_run(...)` in place of the benchmark and the
+launch — and a count the library would read back from the device is taken from the operation's own contract instead
+(`index_wrappers.nonzero_1d_wrapper`). This is how an application enumerates the kernel keys a model forms without a
+device (NeuroBrix: its certification census).
+
+With no observer installed, or one whose `active()` answers False, nothing changes: every launch runs. The observer is
+read at call time, never cached, so one installed after the kernels were defined is the one consulted.
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+__all__ = ["set_launch_observer", "launch_observer"]
+
+
+class _NoObserver:
+    """No observer: never dry, records nothing."""
+
+    @staticmethod
+    def active() -> bool:
+        return False
+
+    @staticmethod
+    def record(tuned, key) -> None:
+        return None                       # every launch is offered to the observer; with none, nothing observes it
+
+    @staticmethod
+    def shadow_run(*args, **kwargs):
+        raise RuntimeError("nbxtensor: a dry run was asked with no launch observer installed")
+
+
+_OBSERVER: Any = _NoObserver
+
+
+def set_launch_observer(observer: Optional[Any]) -> None:
+    """Install `observer` (an object with `active()`, `record(tuned, key)` and `shadow_run(*args, **kwargs)`), or remove
+    it (None). Installing one with a missing method is refused by name."""
+    global _OBSERVER
+    if observer is None:
+        _OBSERVER = _NoObserver
+        return
+    missing = [m for m in ("active", "record", "shadow_run") if not callable(getattr(observer, m, None))]
+    if missing:
+        raise TypeError(f"set_launch_observer: the observer has no callable {', '.join(missing)}")
+    _OBSERVER = observer
+
+
+def launch_observer():
+    """The observer in force (the no-observer when none is installed)."""
+    return _OBSERVER

nbxtensor/ops/layernorm.py

@@ -0,0 +1,71 @@
+"""LayerNorm — pure @triton.jit kernel. Ported from attorch (MIT)."""
+
+import triton
+import triton.language as tl
+from triton import next_power_of_2
+
+from nbxtensor.ops._configs import nbx_autotune, row_program_configs, prune_row_programs
+
+
+# rows per program and warps: the certified configuration of the key (workstream 1), the key the derived census forms
+# (`launch_keys.layer_norm_launches`)
+@nbx_autotune(configs=row_program_configs(), key=['BATCH_BUCKET', 'FEAT_BUCKET', 'scale_by_weight', 'add_bias', 'save_stats'],
+              prune_configs_by={'early_config_prune': prune_row_programs}, cache_results=True)
+@triton.heuristics({
+    'BLOCK_SIZE_FEAT': lambda args: next_power_of_2(args['feat_dim']),
+})
+@triton.jit
+def layer_norm_forward_kernel(
+    input_ptr, weight_ptr, bias_ptr,
+    mean_ptr, inv_std_ptr, output_ptr,
+    batch_dim, feat_dim, BATCH_BUCKET, FEAT_BUCKET,
+    input_batch_stride, input_feat_stride,
+    output_batch_stride, output_feat_stride,
+    eps,
+    scale_by_weight: tl.constexpr,
+    add_bias: tl.constexpr,
+    save_stats: tl.constexpr,
+    BLOCK_SIZE_BATCH: tl.constexpr,
+    BLOCK_SIZE_FEAT: tl.constexpr,
+):
+    """Layer normalization forward.
+
+    input: [batch_dim, feat_dim]
+    weight: [feat_dim] (optional)
+    bias: [feat_dim] (optional)
+    output: [batch_dim, feat_dim]
+    mean, inv_std: [batch_dim] (optional, for backward)
+    """
+    batch_pid = tl.program_id(0).to(tl.int64)
+
+    batch_offset = batch_pid * BLOCK_SIZE_BATCH + tl.arange(0, BLOCK_SIZE_BATCH)
+    feat_offset = tl.arange(0, BLOCK_SIZE_FEAT)
+
+    batch_mask = batch_offset < batch_dim
+    feat_mask = feat_offset < feat_dim
+
+    input_ptr += (input_batch_stride * batch_offset[:, None] +
+                  input_feat_stride * feat_offset[None, :])
+    output_ptr += (output_batch_stride * batch_offset[:, None] +
+                   output_feat_stride * feat_offset[None, :])
+
+    inp = tl.load(input_ptr,
+                  mask=batch_mask[:, None] & feat_mask[None, :]).to(tl.float32)
+    mean = tl.sum(inp, axis=1) / feat_dim
+    diff = tl.where(feat_mask[None, :], inp - mean[:, None], 0)
+    inv_std = tl.rsqrt(tl.sum(diff * diff, axis=1) / feat_dim + eps)
+
+    if save_stats:
+        tl.store(mean_ptr + batch_offset, mean, mask=batch_mask)
+        tl.store(inv_std_ptr + batch_offset, inv_std, mask=batch_mask)
+
+    output = diff * inv_std[:, None]
+    if scale_by_weight:
+        weight = tl.load(weight_ptr + feat_offset, mask=feat_mask)
+        output *= weight
+        if add_bias:
+            bias = tl.load(bias_ptr + feat_offset, mask=feat_mask)
+            output += bias
+
+    tl.store(output_ptr, output,
+             mask=batch_mask[:, None] & feat_mask[None, :])

nbxtensor/ops/le.py

@@ -0,0 +1,35 @@
+"""Less than or equal — pure @triton.jit kernel.
+
+Logic from FlagGems le.py. Both tensor-tensor and scalar variants.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def le_forward_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    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
+
+    x = tl.load(x_ptr + offset, mask=mask)
+    y = tl.load(y_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, (x.to(tl.float32) <= y.to(tl.float32)).to(tl.int1), mask=mask)
+
+@triton.jit
+def le_scalar_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    BLOCK_SIZE: tl.constexpr,
+):
+    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
+
+    x = tl.load(x_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, (x.to(tl.float32) <= scalar).to(tl.int1), mask=mask)

nbxtensor/ops/leaky_relu.py

@@ -0,0 +1,19 @@
+"""Leaky ReLU — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def leaky_relu_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    negative_slope,
+    BLOCK_SIZE: tl.constexpr,
+):
+    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
+
+    x = tl.load(input_ptr + offset, mask=mask)
+    out = tl.maximum(0, x) + negative_slope * tl.minimum(0, x)
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/lerp.py

@@ -0,0 +1,76 @@
+"""Lerp (linear interpolation) kernel — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+
+Tensor variant: result = where(|w| < 0.5, a + w*(b-a), b - (b-a)*(1-w))
+Scalar variant (head): result = a + w*(b-a)
+Scalar variant (tail): result = b - (b-a)*(1-w)
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def lerp_tensor_forward_kernel(
+    input_ptr,
+    end_ptr,
+    weight_ptr,
+    output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Lerp with tensor weight: numerically stable two-branch formula."""
+    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
+
+    a = tl.load(input_ptr + offset, mask=mask)
+    b = tl.load(end_ptr + offset, mask=mask)
+    w = tl.load(weight_ptr + offset, mask=mask)
+
+    result = tl.where(
+        tl.abs(w) < 0.5,
+        a + w * (b - a),
+        b - (b - a) * (1 - w),
+    )
+    tl.store(output_ptr + offset, result, mask=mask)
+
+@triton.jit
+def lerp_scalar_head_kernel(
+    input_ptr,
+    end_ptr,
+    output_ptr,
+    n_elements,
+    weight,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Lerp with scalar weight < 0.5: a + w*(b-a)."""
+    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
+
+    a = tl.load(input_ptr + offset, mask=mask)
+    b = tl.load(end_ptr + offset, mask=mask)
+
+    result = a + weight * (b - a)
+    tl.store(output_ptr + offset, result, mask=mask)
+
+@triton.jit
+def lerp_scalar_tail_kernel(
+    input_ptr,
+    end_ptr,
+    output_ptr,
+    n_elements,
+    weight,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Lerp with scalar weight >= 0.5: b - (b-a)*(1-w)."""
+    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
+
+    a = tl.load(input_ptr + offset, mask=mask)
+    b = tl.load(end_ptr + offset, mask=mask)
+
+    result = b - (b - a) * (1 - weight)
+    tl.store(output_ptr + offset, result, mask=mask)