nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/_cost.py

@@ -0,0 +1,231 @@
+"""The cost an op declares: what it allocates, for how long, what it moves and computes (R-NBXT 4) — prototype.
+
+A planner (NeuroBrix's Prism) prices a model from these declarations instead of re-deriving what the allocator and the
+kernels do. Everything here is plain data (R-NBXT 11): an operand is a `Spec(shape, strides, dtype)`, a cost is a
+`Cost` of `Buffer`s with integer lifetimes, byte counts and flop counts. The declarations read the same layout rules
+the ops execute (`nbxtensor.shape`) and the same widths (`nbxtensor.dtype`).
+
+Prototype scope (docs/nbxtensor/COST.md): the element-wise, reduction and normalisation families, mirroring today's
+wrappers (NeuroBrix `kernels/wrappers.py`) line by line, each declaration citing the allocation it mirrors. Where a
+wrapper takes a decision the library does not hold yet (the binary ops' result dtype, the DtypeEngine's output
+dtype), the caller passes it in; when the wrappers move into the library (M2/M3), `cost` reads the same function the
+kernel launch reads, so the declaration and the execution cannot drift.
+
+Lifetimes are steps of the op's own sequence: a buffer lives over [first, last]; today every transient lives to the
+op's return (the wrappers keep it in a local name); the operands are the caller's and are never counted. `peak_bytes` is the most the op holds at once beyond its operands, outputs included.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Sequence, Tuple
+
+from nbxtensor.dtype.itemsize import itemsize
+from nbxtensor.shape import contiguous_strides, is_contiguous, numel
+
+__all__ = ["Spec", "Buffer", "Cost", "elementwise_unary", "elementwise_binary", "softmax", "rms_norm",
+           "native_layer_norm", "reduction"]
+
+
+@dataclass(frozen=True)
+class Spec:
+    """An operand: its shape, its strides in elements, its dtype name."""
+    shape: Tuple[int, ...]
+    strides: Tuple[int, ...]
+    dtype: str
+
+    @staticmethod
+    def contiguous(shape: Sequence[int], dtype: str) -> "Spec":
+        return Spec(tuple(shape), contiguous_strides(tuple(shape)), dtype)
+
+    @property
+    def nbytes(self) -> int:
+        return numel(self.shape) * itemsize(self.dtype)
+
+
+@dataclass(frozen=True)
+class Buffer:
+    """A buffer the op allocates: its bytes, what it is, and the steps it lives over (`first` to `last`, inclusive).
+    `returned` buffers are the op's outputs: they outlive the op (their `last` is the op's last step)."""
+    name: str
+    nbytes: int
+    first: int
+    last: int
+    returned: bool = False
+
+
+@dataclass(frozen=True)
+class Cost:
+    buffers: Tuple[Buffer, ...]
+    bytes_moved: int        # every byte read and written, copies included
+    flops: int
+
+    @property
+    def outputs(self) -> Tuple[Buffer, ...]:
+        return tuple(b for b in self.buffers if b.returned)
+
+    @property
+    def workspace(self) -> Tuple[Buffer, ...]:
+        return tuple(b for b in self.buffers if not b.returned)
+
+    @property
+    def peak_bytes(self) -> int:
+        steps = {s for b in self.buffers for s in (b.first, b.last)}
+        return max((sum(b.nbytes for b in self.buffers if b.first <= s <= b.last) for s in steps), default=0)
+
+
+class _Steps:
+    """The op's sequence: each allocation takes the next step. Every buffer lives until the op returns: the wrappers
+    hold their transients in local names (the cast, the contiguous copy, the accumulator), which Python releases at
+    the return, never earlier."""
+
+    def __init__(self):
+        self.buffers = []
+        self.moved = 0
+        self.step = 0
+
+    def alloc(self, name: str, nbytes: int, returned: bool = False) -> int:
+        self.step += 1
+        self.buffers.append([name, int(nbytes), self.step, None, returned])
+        return len(self.buffers) - 1
+
+    def free(self, index: int) -> None:
+        """The buffer dies here (a transient the wrapper drops before it returns: a reassigned local, a callee's
+        return): it lives over [its step, this step]."""
+        self.step += 1
+        self.buffers[index][3] = self.step
+
+    def done(self, flops: int) -> Cost:
+        end = self.step
+        return Cost(tuple(Buffer(n, b, f, end if l is None else l, r) for n, b, f, l, r in self.buffers),
+                    self.moved, int(flops))
+
+    def nest(self, cost: "Cost", prefix: str) -> list:
+        """A callee's cost inside this op's sequence: its buffers allocated in its order, its transients freed when it
+        returns; its returned buffers are the caller's to free (`free` on the indices returned)."""
+        base = self.step
+        kept = []
+        for b in cost.buffers:
+            self.buffers.append([f"{prefix}: {b.name}", b.nbytes, base + b.first,
+                                 None if b.returned else base + b.last, False])
+            if b.returned:
+                kept.append(len(self.buffers) - 1)
+        self.step = base + max((b.last for b in cost.buffers), default=0)
+        self.moved += cost.bytes_moved
+        return kept
+
+
+def _contiguous_copy(st: _Steps, x: Spec, name: str) -> Optional[int]:
+    """`x.contiguous()`: no buffer when `x` already lies row-major (`shape.is_contiguous`), else one copy of its
+    elements, read strided and written dense."""
+    if is_contiguous(x.shape, x.strides):
+        return None
+    st.moved += 2 * x.nbytes
+    return st.alloc(name, x.nbytes)
+
+
+def elementwise_unary(x: Spec, out_dtype: Optional[str] = None, flops_per_element: int = 1) -> Cost:
+    """A unary element-wise op (`wrappers._prepare_unary`: `x.contiguous()`, then `empty_like` the output; `out_dtype`
+    when the op writes another dtype, e.g. a comparison's bool)."""
+    st = _Steps()
+    _contiguous_copy(st, x, "contiguous input")
+    out = Spec.contiguous(x.shape, out_dtype or x.dtype)
+    st.alloc("output", out.nbytes, returned=True)
+    st.moved += x.nbytes + out.nbytes
+    return st.done(flops_per_element * numel(x.shape))
+
+
+def elementwise_binary(a: Spec, b: Spec, common_dtype: str, out_shape: Sequence[int],
+                       out_dtype: Optional[str] = None, flops_per_element: int = 1) -> Cost:
+    """A binary element-wise op of two tensors (`wrappers._prepare_binary_strided`): each operand not already at
+    `common_dtype` is cast — a COPY at that dtype (the RoPE lesson: an in-kernel widen moves the compiler's
+    contraction) —, then both are read by strides as broadcast VIEWS (no expand copy), and the output is dense at
+    the broadcast shape. `common_dtype` and `out_shape` are the wrapper's decisions, passed in until the wrapper's
+    promotion lives here."""
+    st = _Steps()
+    for name, x in (("a cast", a), ("b cast", b)):
+        if x.dtype != common_dtype:
+            cast = Spec(x.shape, x.strides, common_dtype)
+            st.moved += x.nbytes + cast.nbytes
+            st.alloc(name, cast.nbytes)
+    out = Spec.contiguous(out_shape, out_dtype or common_dtype)
+    st.alloc("output", out.nbytes, returned=True)
+    st.moved += numel(a.shape) * itemsize(common_dtype) + numel(b.shape) * itemsize(common_dtype) + out.nbytes
+    return st.done(flops_per_element * numel(out_shape))
+
+
+def softmax(x: Spec, dim: int) -> Cost:
+    """`wrappers.softmax`: the reduced dim moved last (`movedim(dim, -1).contiguous()`: a copy unless that view is
+    already dense), the output `empty_like` the 2-D view, returned moved back (a view). Fused and looped kernels
+    alike hold nothing else. Arithmetic per element: the row max, the shift, the exponential, the row sum, the
+    division."""
+    dim = dim % len(x.shape)
+    order = [d for d in range(len(x.shape)) if d != dim] + [dim]
+    moved = Spec(tuple(x.shape[d] for d in order), tuple(x.strides[d] for d in order), x.dtype)
+    st = _Steps()
+    _contiguous_copy(st, moved, "rows made contiguous")
+    st.alloc("output", x.nbytes, returned=True)
+    st.moved += x.nbytes + x.nbytes
+    return st.done(5 * numel(x.shape))
+
+
+def rms_norm(x: Spec, has_weight: bool, out_dtype: Optional[str] = None) -> Cost:
+    """`wrappers.rms_norm`: `x.contiguous()`; when that made a fresh copy and the output dtype is the input's, the
+    output is written INTO the copy (no second buffer); otherwise a new output. Arithmetic per element: the square,
+    the row sum, the scale by the reciprocal root, and the weight's product when there is one."""
+    out_dt = out_dtype or x.dtype
+    st = _Steps()
+    copy = _contiguous_copy(st, x, "contiguous input")
+    out_bytes = numel(x.shape) * itemsize(out_dt)
+    if copy is not None and out_dt == x.dtype:
+        st.buffers[copy][4] = True                                 # the copy IS the output
+        st.buffers[copy][0] = "output (in the contiguous copy)"
+    else:
+        st.alloc("output", out_bytes, returned=True)
+    st.moved += x.nbytes + out_bytes + (x.shape[-1] * itemsize(x.dtype) if has_weight else 0)
+    return st.done((4 if has_weight else 3) * numel(x.shape))
+
+
+def native_layer_norm(x: Spec, normalized_shape: Sequence[int], has_weight: bool, has_bias: bool) -> Cost:
+    """`wrappers.native_layer_norm`: `x.contiguous()`, the output `empty_like` it, and the row statistics (mean and
+    1/std, float32, one per row) — returned, as aten::native_layer_norm returns (output, mean, rstd). Arithmetic per
+    element: the mean's sum, the centring, the square, the variance's sum, the scale; the weight's product and the
+    bias's sum when present."""
+    feat = numel(normalized_shape)
+    rows = numel(x.shape) // feat
+    st = _Steps()
+    _contiguous_copy(st, x, "contiguous input")
+    st.alloc("output", x.nbytes, returned=True)
+    st.alloc("mean", rows * itemsize("float32"), returned=True)
+    st.alloc("rstd", rows * itemsize("float32"), returned=True)
+    affine = feat * itemsize(x.dtype) * (int(has_weight) + int(has_weight and has_bias))
+    st.moved += 2 * x.nbytes + 2 * rows * itemsize("float32") + affine
+    return st.done((5 + int(has_weight) + int(has_weight and has_bias)) * numel(x.shape))
+
+
+def reduction(x: Spec, dim: Optional[int], keepdim: bool, out_dtype: str, cast_to: Optional[str] = None) -> Cost:
+    """`wrappers.sum_wrapper` / `mean_wrapper` over one dim or all: the input cast to the op's `dtype` keyword when
+    given (`cast_to`, a copy), the reduced dim moved last and made contiguous (a copy unless already dense), one
+    float32 accumulator per output element, then cast to the result dtype (`out_dtype`: the input's for a float
+    input, int64 for sum of an integer one) — a second buffer unless that dtype is float32 (`.to` returns the tensor
+    itself then). One addition per input element."""
+    st = _Steps()
+    src = x
+    if cast_to is not None and cast_to != x.dtype:
+        src = Spec(x.shape, x.strides, cast_to)
+        st.moved += x.nbytes + src.nbytes
+        st.alloc("dtype keyword cast", src.nbytes)
+    if dim is None:
+        rows = 1
+        _contiguous_copy(st, src, "flattened input")
+    else:
+        d = dim % len(src.shape)
+        order = [i for i in range(len(src.shape)) if i != d] + [d]
+        moved = Spec(tuple(src.shape[i] for i in order), tuple(src.strides[i] for i in order), src.dtype)
+        rows = numel(src.shape) // src.shape[d]
+        _contiguous_copy(st, moved, "reduced dim made contiguous")
+    st.alloc("float32 accumulator", rows * itemsize("float32"), returned=out_dtype == "float32")
+    st.moved += src.nbytes + rows * itemsize("float32")
+    if out_dtype != "float32":
+        st.alloc("result", rows * itemsize(out_dtype), returned=True)
+        st.moved += rows * itemsize("float32") + rows * itemsize(out_dtype)
+    return st.done(numel(x.shape))

nbxtensor/ops/abs.py

@@ -0,0 +1,17 @@
+"""Absolute value — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def abs_forward_kernel(
+    input_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(input_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, tl.abs(x), mask=mask)

nbxtensor/ops/adaptive_avg_pool2d.py

@@ -0,0 +1,161 @@
+"""Adaptive average pooling 2D — pure @triton.jit kernel.
+
+For each output element (n, c, oh, ow), the kernel computes the input
+window boundaries using the standard PyTorch adaptive pooling formula:
+    start_h = floor(oh * IH / OH)
+    end_h   = ceil((oh + 1) * IH / OH)
+and averages all input elements in [start_h, end_h) x [start_w, end_w).
+
+The end index is a CEILING. It was a floor here, which is the same value
+whenever the output size divides the input size and a smaller window
+whenever it does not — so every exactly-dividing shape agreed and every
+ragged one was quietly averaging too few elements. Measured against torch
+2026-09-09 by one-hot probing which input rows reach each output row:
+IH=7,OH=3 gives torch [(0,3),(2,5),(4,7)] and the floor form
+[(0,2),(2,4),(4,7)].
+
+The window loop was also fixed at 8x8, on a comment that the adaptive
+window is "typically 1-3 pixels each dim". That is false for any pooling
+that downsamples by more than 8, and global average pooling is the common
+case: a 16x16 plane to 1x1 has a 16x16 window, of which the kernel summed
+the first 8x8 and divided by 64. Measured: 0.1807 returned where the plane
+mean is 0.0019. The bound is now derived from the shape.
+
+BOTH defects are in this kernel, not in any backend, so both were wrong on
+EVERY vendor. This changes NVIDIA's numbers, and is meant to: keeping them
+would be keeping the bug.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def adaptive_avg_pool2d_kernel(
+    input_ptr,
+    output_ptr,
+    n_elements,
+    # Input shape: [N, C, IH, IW]
+    IH, IW,
+    # Output shape: [N, C, OH, OW]
+    OH, OW,
+    # Strides (input contiguous NCHW)
+    stride_n, stride_c, stride_h, stride_w,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Adaptive average pool 2D.
+
+    Each thread computes one output element by averaging over the
+    corresponding adaptive window in the input.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    idx = 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 = idx < n_elements
+
+    # Decompose flat index -> (n, c, oh, ow)
+    ow = idx % OW
+    tmp = idx // OW
+    oh = tmp % OH
+    tmp2 = tmp // OH
+    c = tmp2 % (stride_n // stride_c)  # C = stride_n / stride_c for contiguous
+    n = tmp2 // (stride_n // stride_c)
+
+    # Adaptive pooling window bounds (PyTorch formula)
+    # start = floor(out_idx * in_size / out_size)
+    # end   = CEIL((out_idx + 1) * in_size / out_size)
+    h_start = (oh * IH) // OH
+    h_end = ((oh + 1) * IH + OH - 1) // OH
+    w_start = (ow * IW) // OW
+    w_end = ((ow + 1) * IW + OW - 1) // OW
+
+    # Base pointer for this (n, c) plane
+    base = n * stride_n + c * stride_c
+
+    # Accumulate sum over the window
+    # Window size varies per output position in adaptive pooling
+    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
+    count = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
+
+    # The widest window this shape can produce, derived rather than assumed.
+    # end - start <= ceil((i+1)*IN/OUT) - floor(i*IN/OUT) < IN/OUT + 2, so
+    # floor(IN/OUT) + 2 bounds it for every i. The old fixed 8 silently
+    # truncated any window wider than that — every global average pool over
+    # a plane larger than 8x8 summed a corner and divided by its own count.
+    KH = IH // OH + 2
+    KW = IW // OW + 2
+    for dh in range(0, KH):
+        for dw in range(0, KW):
+            h_idx = h_start + dh
+            w_idx = w_start + dw
+            in_bounds = (h_idx < h_end) & (w_idx < w_end)
+            valid = mask & in_bounds
+
+            offset = base + h_idx * stride_h + w_idx * stride_w
+            val = tl.load(input_ptr + offset, mask=valid, other=0.0).to(tl.float32)
+            acc += tl.where(valid, val, 0.0)
+            count += tl.where(valid, 1.0, 0.0)
+
+    # Average
+    count = tl.maximum(count, 1.0)  # Avoid division by zero
+    result = acc / count
+
+    # Store
+    out_offset = idx
+    tl.store(output_ptr + out_offset, result, mask=mask)
+
+@triton.jit
+def adaptive_avg_pool2d_tiled_kernel(
+    input_ptr,
+    output_ptr,
+    N, C,
+    IH, IW,
+    OH, OW,
+    BLOCK_X: tl.constexpr,
+    BLOCK_Y: tl.constexpr,
+):
+    """Tiled variant for spatial parallelism.
+
+    Grid: (cdiv(OW, BLOCK_X), cdiv(OH, BLOCK_Y), N*C)
+    """
+    pid_x = tl.program_id(0).to(tl.int64)
+    pid_y = tl.program_id(1).to(tl.int64)
+    pid_nc = tl.program_id(2).to(tl.int64)
+
+    ow = pid_x * BLOCK_X + tl.arange(0, BLOCK_X)  # [BLOCK_X]
+    oh = pid_y * BLOCK_Y + tl.arange(0, BLOCK_Y)  # [BLOCK_Y]
+
+    mask_h = oh < OH
+    mask_w = ow < OW
+    mask_2d = mask_h[:, None] & mask_w[None, :]
+
+    # Window boundaries
+    h_start = (oh * IH) // OH  # [BLOCK_Y]
+    h_end = ((oh + 1) * IH) // OH
+    w_start = (ow * IW) // OW  # [BLOCK_X]
+    w_end = ((ow + 1) * IW) // OW
+
+    # Window sizes
+    kh = h_end - h_start  # [BLOCK_Y]
+    kw = w_end - w_start  # [BLOCK_X]
+
+    base = pid_nc * IH * IW
+
+    acc = tl.zeros([BLOCK_Y, BLOCK_X], dtype=tl.float32)
+
+    for dh in range(0, 4):
+        for dw in range(0, 4):
+            h_idx = h_start[:, None] + dh
+            w_idx = w_start[None, :] + dw
+            in_bounds = (dh < kh[:, None]) & (dw < kw[None, :])
+            valid = mask_2d & in_bounds
+
+            offset = base + h_idx * IW + w_idx
+            val = tl.load(input_ptr + offset, mask=valid, other=0.0).to(tl.float32)
+            acc += tl.where(valid, val, 0.0)
+
+    area = (kh[:, None] * kw[None, :]).to(tl.float32)
+    area = tl.maximum(area, 1.0)
+    result = acc / area
+
+    out_base = pid_nc * OH * OW
+    out_off = out_base + oh[:, None] * OW + ow[None, :]
+    tl.store(output_ptr + out_off, result, mask=mask_2d)

nbxtensor/ops/add.py

@@ -0,0 +1,96 @@
+"""Add — pure @triton.jit kernel (tensor + tensor, tensor + scalar)."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def add_forward_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    alpha,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x + alpha * y (tensor + tensor)"""
+    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 + alpha * y, mask=mask)
+
+@triton.jit
+def add_scalar_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x + scalar"""
+    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 + scalar, mask=mask)
+
+
+@triton.jit
+def add_scalar_dev_kernel(
+    x_ptr, s_ptr, output_ptr,
+    n_elements,
+    alpha,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x + (s * alpha) where s is a 0-d DEVICE tensor.
+
+    Bit-exact mirror of add_scalar_kernel's `x + scalar` path: the
+    host path computes `float(s) * alpha` in Python float64 then
+    narrows to the f32 kernel argument; here the same math runs
+    in-register (load -> f64 -> *alpha -> f32). Replaces the host
+    .item() sync so the scalar stays device-resident — the replay
+    device-scalar increment (Ming timestep class, 2026-08-15).
+
+    Bit-exactness constraint: `alpha` binds as an f32 runtime arg, so
+    the f64 multiply sees f32(alpha) where the host path sees the full
+    f64 alpha — a 1-ulp double-rounding risk for alpha values not
+    exactly representable in f32. Every ATen call site passes alpha=±1
+    (exact); a non-representable alpha would need an f64-carried arg.
+    """
+    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
+    s = (tl.load(s_ptr).to(tl.float64) * alpha).to(tl.float32)
+    x = tl.load(x_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, x + s, mask=mask)
+
+
+@triton.jit
+def add_bias_broadcast_kernel(
+    x_ptr, bias_ptr, output_ptr,
+    n_elements, feat_dim,
+    alpha,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out[i] = x[i] + alpha * bias[i % feat_dim].
+
+    Used by `add` wrapper when `y` is a 1D bias of size `feat_dim` that
+    broadcasts against `x`'s last dim. Avoids materializing the 8 GiB
+    contiguous broadcast view of bias on Sana 4Kpx VAE add::88
+    (`mul::58::out_0 (1, 4096, 4096, 128) + bias (128)`). Reads bias
+    directly via `offset % feat_dim` indexing — same effect as a
+    stride-0 broadcast but compatible with the kernel's flat 1D
+    addressing.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    # Cast to int64 — tensors >= 2^31 elements (e.g. Sana 4Kpx VAE
+    # add::88 input 1x4096x4096x128 = 2^31) overflow int32 offset
+    # arithmetic silently and corrupt `offset % feat_dim`, producing
+    # garbage output. P-SANA-4KPX-RUNTIME 2026-05-07.
+    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
+
+    c = offset % feat_dim
+    x = tl.load(x_ptr + offset, mask=mask)
+    b = tl.load(bias_ptr + c, mask=mask)
+    tl.store(output_ptr + offset, x + alpha * b, mask=mask)

nbxtensor/ops/addcdiv.py

@@ -0,0 +1,29 @@
+"""Addcdiv kernel — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+Logic: result = input + value * (tensor1 / tensor2)
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def addcdiv_forward_kernel(
+    input_ptr,
+    tensor1_ptr,
+    tensor2_ptr,
+    output_ptr,
+    n_elements,
+    value,
+    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)
+    t1 = tl.load(tensor1_ptr + offset, mask=mask)
+    t2 = tl.load(tensor2_ptr + offset, mask=mask)
+
+    result = x + value * (t1 / t2)
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/addcmul.py

@@ -0,0 +1,29 @@
+"""Addcmul kernel — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+Logic: result = input + value * tensor1 * tensor2
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def addcmul_forward_kernel(
+    input_ptr,
+    tensor1_ptr,
+    tensor2_ptr,
+    output_ptr,
+    n_elements,
+    value,
+    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)
+    t1 = tl.load(tensor1_ptr + offset, mask=mask)
+    t2 = tl.load(tensor2_ptr + offset, mask=mask)
+
+    result = x + value * t1 * t2
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/addmv_op.py

@@ -0,0 +1,64 @@
+"""Addmv — pure @triton.jit kernel.
+
+Ported from FlagGems addmv (Apache-2.0 license).
+Computes: out = beta * self + alpha * (mat @ vec)
+where mat is [N, M], vec is [M], self is broadcastable to [N].
+"""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops._configs import nbx_autotune, addmv_configs, prune_addmv
+
+
+# rows, reduction tile, warps: the certified configuration of the key (workstream 1, small-M GEMM), the key the derived
+# census forms (`launch_keys.addmv_launches`)
+@nbx_autotune(configs=addmv_configs(), key=['N_BUCKET', 'M_BUCKET'],
+              prune_configs_by={'early_config_prune': prune_addmv}, cache_results=True)
+@triton.jit
+def addmv_kernel(
+    A_ptr,
+    B_ptr,
+    inp_ptr,
+    out_ptr,
+    N, M, N_BUCKET, M_BUCKET,
+    alpha,
+    beta,
+    stride_an, stride_am,
+    stride_bm,
+    stride_in,
+    stride_outn,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Matrix-vector multiply with bias: out = beta*inp + alpha*(A @ B).
+
+    Each program handles BLOCK_N rows of A, iterating over M in BLOCK_M tiles.
+    Accumulates in fp32 for stability.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    offset_n = pid * BLOCK_N + tl.arange(0, BLOCK_N)[:, None]
+    offset_m = tl.arange(0, BLOCK_M)[None, :]
+    n_mask = offset_n < N
+
+    A_ptrs = A_ptr + offset_n * stride_an + offset_m * stride_am
+    B_ptrs = B_ptr + offset_m * stride_bm
+
+    acc = tl.zeros((BLOCK_N, BLOCK_M), dtype=tl.float32)
+    for m in range(0, M, BLOCK_M):
+        m_mask = m + offset_m < M
+        a = tl.load(A_ptrs, mask=n_mask & m_mask, other=0.0).to(tl.float32)
+        b = tl.load(B_ptrs, mask=m_mask, other=0.0).to(tl.float32)
+        acc += a * b
+        A_ptrs += BLOCK_M * stride_am
+        B_ptrs += BLOCK_M * stride_bm
+
+    # Reduce over M dimension, then add bias
+    mv_result = tl.sum(acc, axis=1)[:, None]
+    inp_ptrs = inp_ptr + offset_n * stride_in
+    inp = tl.load(inp_ptrs, mask=n_mask, other=0.0).to(tl.float32)
+
+    out = mv_result * alpha + inp * beta
+
+    out_ptrs = out_ptr + offset_n * stride_outn
+    tl.store(out_ptrs, out, mask=n_mask)