nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/div.py

@@ -0,0 +1,75 @@
+"""Division — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def div_forward_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x / 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 / y, mask=mask)
+
+@triton.jit
+def div_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 div_scalar_dev_kernel(
+    x_ptr, s_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x / s where s is a 0-d DEVICE tensor. Bit-exact mirror
+    of div_scalar_kernel (host float(s) -> f32 arg == load -> f64 ->
+    f32). Device-scalar increment 2026-08-15."""
+    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).to(tl.float32)
+    x = tl.load(x_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, x / s, mask=mask)
+
+
+@triton.jit
+def int_div_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    FLOOR: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x // y on integers, computed in the integers: truncated toward zero (C), or toward
+    -inf when FLOOR (ATen's floor_divide). Exact at every magnitude of the dtype, where a float
+    round-trip is not: Metal's float32 divide is not correctly rounded (-3403868 / 13 gave
+    -261836.02, floor -261837; 4435 of 10000 quotients differ from IEEE, 2026-10-08)."""
+    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, other=0)
+    y = tl.load(y_ptr + offset, mask=mask, other=1)
+    q = x // y
+    if FLOOR:
+        r = x - q * y
+        q = q - ((r != 0) & ((r < 0) != (y < 0))).to(q.dtype)
+    tl.store(output_ptr + offset, q, mask=mask)

nbxtensor/ops/dot_op.py

@@ -0,0 +1,67 @@
+"""Dot product — pure @triton.jit kernels.
+
+Ported from FlagGems dot (Apache-2.0 license).
+Two-phase reduction for large N: partial sums then final reduce.
+Single-phase for small N.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def dot_kernel_small(
+    x_ptr, y_ptr, out_ptr,
+    N,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Single-block dot product for small vectors (N < 4096).
+
+    Accumulates in fp32 for numerical stability.
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < N
+
+    x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+    y = tl.load(y_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+    result = tl.sum(x * y)
+    tl.store(out_ptr, result)
+
+
+@triton.jit
+def dot_kernel_partial(
+    x_ptr, y_ptr, mid_ptr,
+    N,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Phase 1: compute partial dot products per block.
+
+    Each block computes sum(x[block] * y[block]) and stores to mid_ptr[pid].
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < N
+
+    x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+    y = tl.load(y_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+    partial_sum = tl.sum(x * y)
+    tl.store(mid_ptr + pid, partial_sum)
+
+
+@triton.jit
+def dot_kernel_reduce(
+    mid_ptr, out_ptr,
+    M,
+    BLOCK_MID: tl.constexpr,
+):
+    """Phase 2: reduce partial sums into final scalar result."""
+    offsets = tl.arange(0, BLOCK_MID)
+    mask = offsets < M
+    mid_val = tl.load(mid_ptr + offsets, mask=mask, other=0.0)
+    result = tl.sum(mid_val)
+    tl.store(out_ptr, result)

nbxtensor/ops/dropout.py

@@ -0,0 +1,23 @@
+"""Dropout — pure @triton.jit kernel.
+
+At inference time (train=False), dropout is identity (passthrough).
+This kernel implements the identity passthrough for inference mode.
+Training dropout with Philox RNG would be a separate kernel.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def dropout_inference_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Dropout at inference = identity copy."""
+    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, mask=mask)

nbxtensor/ops/dtype_convert.py

@@ -0,0 +1,42 @@
+"""GPU dtype conversion kernels — bf16↔fp16 via bit manipulation.
+
+Eliminates CPU-side numpy conversion during weight loading.
+Raw bf16 bytes are H2D copied as uint16, then converted on GPU.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def bf16_to_fp16_kernel(
+    src_ptr,   # uint16 (bf16 raw bits) on GPU
+    dst_ptr,   # fp16 output on GPU
+    N,         # total number of elements
+    BLOCK: tl.constexpr,
+):
+    """Convert bf16 → fp16 on GPU via fp32 intermediate.
+
+    bf16 bits:  [sign(1)][exp(8)][mantissa(7)]
+    fp32 bits:  [sign(1)][exp(8)][mantissa(23)]  ← bf16 << 16
+    fp16 bits:  [sign(1)][exp(5)][mantissa(10)]  ← GPU truncates fp32→fp16
+
+    The GPU handles fp32→fp16 truncation natively (with rounding).
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    offs = pid.to(tl.int64) * BLOCK + tl.arange(0, BLOCK)   # 64-bit from the program id (register 58)
+    mask = offs < N
+
+    # Load bf16 raw bits (stored as int16, same 2-byte layout)
+    raw_i16 = tl.load(src_ptr + offs, mask=mask, other=0)
+    # Mask to 16 bits to prevent sign extension when widening to uint32
+    raw = (raw_i16.to(tl.int32) & 0xFFFF).to(tl.uint32)
+
+    # bf16 → fp32: shift left 16 bits, bitcast to float32
+    fp32_bits = raw << 16
+    fp32_vals = fp32_bits.to(tl.float32, bitcast=True)
+
+    # fp32 → fp16: GPU truncates natively
+    fp16_vals = fp32_vals.to(tl.float16)
+
+    tl.store(dst_ptr + offs, fp16_vals, mask=mask)

nbxtensor/ops/elementwise_launch.py

@@ -0,0 +1,81 @@
+"""The launch configuration of every element-wise kernel, certified (chantier 1).
+
+An element-wise kernel reads and writes each element once: its time is the memory's, and the program's tile and warps
+move it. Measured on V100 (2026-10-23, `campaigns/2026_10_12_mq25_gate/stable/c1rest/ew_bench.json`: add, mul, exp,
+silu, gelu, erf, pow and a copy, fp32 and fp16, 4 096 to 16 777 216 elements, 15 tiles x warps): at 16 M elements the
+best of the 15 is within 0.3 % (max 1.4 %) for every operation of the configuration that is best for the copy; below
+65 536 elements the launch itself dominates. The OPERATION does not move the optimum — the element count does, by
+class. So the element-wise kernels share ONE key: the class of the launch's element count (`autotune.buckets.elements`,
+the profile's ladder, else its `default`), certified on the copy below, and every element-wise launch reads its tile
+and warps through `ew_meta` — one memoized resolution per class and process, through the autotuner's own door
+(`ops/_configs.nbx_autotune`: the census record, the certified directory, the refusal of a confirmation run).
+
+The copy is in fp32 (the widest element a launch of this class commonly moves); a tile changes no bit of any
+element-wise result — each element is computed alone — so its oracle is identity.
+"""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops._configs import nbx_autotune
+
+
+def _ew_configs():
+    """Tile x warps; the configuration every element-wise launch ran before certification (1 024 x 4) is among them."""
+    return [triton.Config({'BLOCK_SIZE': block}, num_warps=warps)
+            for block in (256, 512, 1024, 2048, 4096) for warps in (2, 4, 8)]
+
+
+@nbx_autotune(configs=_ew_configs(), key=['N_BUCKET'], cache_results=True)
+@triton.jit
+def ew_launch_kernel(x_ptr, y_ptr, n_elements, N_BUCKET, BLOCK_SIZE: tl.constexpr):
+    """y = x over n_elements: the element-wise class's representative."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offset < n_elements
+    tl.store(y_ptr + offset, tl.load(x_ptr + offset, mask=mask), mask=mask)
+
+
+_MEMO = {}          # {the element count's class: {'BLOCK_SIZE': ..., 'num_warps': ...}} — one resolution per process
+
+
+def ew_bucket(n_elements: int) -> int:
+    """The class of an element-wise launch of `n_elements` (`autotune.buckets.elements`, else the profile's
+    `default`): the key `launch_keys.elementwise_launches` forms for the census."""
+    from nbxtensor.ops.autotune_bucket import bucket_of
+    return bucket_of("elements", max(1, int(n_elements)))
+
+
+def ew_meta(n_elements: int) -> dict:
+    """The tile and warps of an element-wise launch of `n_elements`: its class's certified configuration (a
+    confirmation run refuses a class the directory lacks, by name; another run sweeps it once), read once per class
+    and process. The resolution launches the copy above at the class's floor — the fewest elements that form it, so
+    never more than the launch that asks — on a scratch buffer of the current device, through the autotuner's door."""
+    nb = ew_bucket(n_elements)
+    hit = _MEMO.get(nb)
+    if hit is not None:
+        return hit
+    from nbxtensor.ops.autotune_bucket import class_floor
+    from nbxtensor.tensor import NBXTensor
+    from nbxtensor.dtype.nbx_dtype import NBXDtype
+    from nbxtensor.ops.launch_observer import launch_observer
+    _census = launch_observer()                 # the application's launch observer (dry run), if any
+    from nbxtensor.ops.autotune_cache import key_of
+    rep = class_floor("elements", nb)
+    buf = NBXTensor.empty((rep,), dtype=NBXDtype.float32, device='cuda')
+    args = (buf, buf, rep, nb)
+    ew_launch_kernel[lambda META: (triton.cdiv(rep, META['BLOCK_SIZE']),)](*args)
+    cfg = ew_launch_kernel.cache.get(key_of(ew_launch_kernel, args, {}))
+    if cfg is None:
+        if _census.active():         # a census shadow launches nothing: no configuration exists, none is needed
+            first = ew_launch_kernel.configs[0]
+            return {'BLOCK_SIZE': first.kwargs['BLOCK_SIZE'], 'num_warps': first.num_warps}
+        raise RuntimeError(f"element-wise class {nb}: the copy's autotuner ran and holds no configuration for it")
+    hit = {'BLOCK_SIZE': int(cfg.kwargs['BLOCK_SIZE']), 'num_warps': int(cfg.num_warps)}
+    _MEMO[nb] = hit
+    return hit
+
+
+def ew_grid(n_elements: int, meta: dict) -> tuple:
+    """The 1-D grid of an element-wise launch at its configuration."""
+    return (triton.cdiv(int(n_elements), meta['BLOCK_SIZE']),)