nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/batch_norm.py

@@ -0,0 +1,163 @@
+"""Batch normalization — pure @triton.jit kernel. Extracted from FlagGems.
+
+For inference: uses running_mean/running_var (is_train=False path).
+For training: computes online mean/var with Welford's algorithm.
+"""
+
+import triton
+import triton.language as tl
+from triton import next_power_of_2
+
+
+def _batch_norm_block_heuristic(args):
+    """Heuristic for BLOCK_M: process multiple batch elements per CTA when
+    spatial_dim is small."""
+    return (min(max(1, next_power_of_2(args['batch_dim'] // 2**10)), 128)
+            if args['spatial_dim'] < 64 else 1)
+
+
+def _batch_norm_block_n_heuristic(args):
+    return min(4096, next_power_of_2(args['spatial_dim']))
+
+
+@triton.heuristics({
+    'BLOCK_M': _batch_norm_block_heuristic,
+    'BLOCK_N': _batch_norm_block_n_heuristic,
+})
+@triton.jit
+def batch_norm_forward_kernel(
+    input_pointer,
+    weight_pointer,
+    bias_pointer,
+    mean_pointer,
+    inv_std_pointer,
+    output_pointer,
+    running_mean_pointer,
+    running_var_pointer,
+    batch_dim,
+    spatial_dim,
+    input_batch_stride,
+    input_feat_stride,
+    input_spatial_stride,
+    output_batch_stride,
+    output_feat_stride,
+    output_spatial_stride,
+    momentum,
+    eps,
+    is_train: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    feat_pid = tl.program_id(0).to(tl.int64)
+
+    # Training mode: compute mean/var online
+    if is_train:
+        mean = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        var = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        cnt = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.int32)
+
+        m_num_steps = tl.cdiv(batch_dim, BLOCK_M)
+        n_num_steps = tl.cdiv(spatial_dim, BLOCK_N)
+
+        for m_step in range(0, m_num_steps):
+            for n_step in range(0, n_num_steps):
+                spatial_offset = n_step * BLOCK_N + tl.arange(0, BLOCK_N)
+                spatial_mask = spatial_offset < spatial_dim
+
+                batch_offset = m_step * BLOCK_M + tl.arange(0, BLOCK_M)
+                batch_mask = batch_offset < batch_dim
+
+                curr_input_pointer = (
+                    input_pointer
+                    + input_feat_stride * feat_pid
+                    + input_batch_stride * batch_offset[:, None]
+                    + input_spatial_stride * spatial_offset[None, :]
+                )
+
+                mask = batch_mask[:, None] & spatial_mask[None, :]
+                curr_input = tl.load(curr_input_pointer, mask=mask).to(tl.float32)
+
+                step = m_step * n_num_steps + n_step + 1
+                new_mean = tl.where(mask, mean + (curr_input - mean) / step, mean)
+                new_var = tl.where(
+                    mask, var + (curr_input - new_mean) * (curr_input - mean), var
+                )
+                cnt += mask.to(tl.int32)
+                mean = new_mean
+                var = new_var
+
+        final_mean = tl.sum(mean * cnt) / (batch_dim * spatial_dim)
+        var = tl.sum(var + cnt * (mean - final_mean) * (mean - final_mean)) / (
+            batch_dim * spatial_dim
+        )
+        inv_std = 1.0 / tl.sqrt(var + eps)
+        mean = final_mean
+
+        tl.store(feat_pid + mean_pointer, mean)
+        tl.store(feat_pid + inv_std_pointer, inv_std)
+
+        # Running-stats update only when running buffers are supplied. The
+        # instance_norm case (AdaIN: training=True, no running mean/var) passes
+        # null running pointers — skip the update so we neither dereference a
+        # null pointer nor write the momentum-blended stats back into the input
+        # tensor (the old x-placeholder corrupted the activations in place).
+        if running_mean_pointer:
+            running_mean_pointer += feat_pid
+            running_var_pointer += feat_pid
+
+            running_mean = tl.load(running_mean_pointer)
+            running_var = tl.load(running_var_pointer)
+
+            n = batch_dim * spatial_dim
+            tl.store(running_mean_pointer, (1 - momentum) * running_mean + momentum * mean)
+            tl.store(
+                running_var_pointer,
+                (1 - momentum) * running_var + momentum * var * n / (n - 1),
+            )
+
+    else:
+        # Inference mode: use running stats
+        mean = tl.load(feat_pid + running_mean_pointer)
+        inv_std = 1.0 / tl.sqrt(tl.load(feat_pid + running_var_pointer) + eps)
+
+    if weight_pointer:
+        weight = tl.load(feat_pid + weight_pointer).to(tl.float32)
+    else:
+        weight = 1.0
+    if bias_pointer:
+        bias = tl.load(feat_pid + bias_pointer).to(tl.float32)
+    else:
+        bias = 0.0
+
+    # Apply normalization: weight * (x - mean) * inv_std + bias
+    for m_step in range(0, tl.cdiv(batch_dim, BLOCK_M)):
+        for n_step in range(0, tl.cdiv(spatial_dim, BLOCK_N)):
+            batch_offset = m_step * BLOCK_M + tl.arange(0, BLOCK_M)
+            batch_mask = batch_offset < batch_dim
+
+            spatial_offset = n_step * BLOCK_N + tl.arange(0, BLOCK_N)
+            spatial_mask = spatial_offset < spatial_dim
+
+            curr_input_pointer = (
+                input_pointer
+                + input_feat_stride * feat_pid
+                + input_batch_stride * batch_offset[:, None]
+                + input_spatial_stride * spatial_offset[None, :]
+            )
+            curr_output_pointer = (
+                output_pointer
+                + output_feat_stride * feat_pid
+                + output_batch_stride * batch_offset[:, None]
+                + output_spatial_stride * spatial_offset[None, :]
+            )
+
+            curr_input = tl.load(
+                curr_input_pointer, mask=batch_mask[:, None] & spatial_mask[None, :]
+            ).to(tl.float32)
+            output = weight * (curr_input - mean) * inv_std + bias
+
+            tl.store(
+                curr_output_pointer,
+                output,
+                mask=batch_mask[:, None] & spatial_mask[None, :],
+            )

nbxtensor/ops/binary_strided.py

@@ -0,0 +1,196 @@
+"""Element-wise add, sub, mul and div over two operands read by their own strides —
+pure @triton.jit, N-D scalable like `strided_copy`.
+
+`_prepare_binary` used to broadcast a smaller operand by `expand(...).
+contiguous()` — a full-size transient per call: a conv bias (32,) over a
+448 × 448 image was 6.4 M elements written and read back per layer
+(351 strided copies per Real-ESRGAN upscale, the copy census of
+2026-09-07); an attention mask (S, S) broadcast over the heads the same.
+These kernels read each operand through its strides (0 on a broadcast
+dim), so the broadcast costs nothing and a transposed operand is read
+in place. The arithmetic is the flat kernels' expression, verbatim
+(`x + alpha * y`, `x * y`), so the bytes are those of the copying path.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def add_strided_nd_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    alpha,
+    shape_ptr,      # GPU pointer → NDIM int64 extents of the output
+    x_stride_ptr,   # GPU pointer → NDIM int64 strides of x (0 on a broadcast dim)
+    y_stride_ptr,   # GPU pointer → NDIM int64 strides of y
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+):
+    """out = x + alpha * y, x and y read by their strides, out contiguous."""
+    pid = tl.program_id(0).to(tl.int64)
+    offsets = 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 = offsets < n_elements
+    remaining = offsets
+    x_off = tl.zeros_like(offsets)
+    y_off = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        xs = tl.load(x_stride_ptr + dim)
+        ys = tl.load(y_stride_ptr + dim)
+        idx = remaining % d
+        remaining = remaining // d
+        x_off = x_off + idx * xs
+        y_off = y_off + idx * ys
+    x = tl.load(x_ptr + x_off, mask=mask)
+    y = tl.load(y_ptr + y_off, mask=mask)
+    tl.store(output_ptr + offsets, x + alpha * y, mask=mask)
+
+
+@triton.jit
+def mul_strided_nd_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    shape_ptr,
+    x_stride_ptr,
+    y_stride_ptr,
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+):
+    """out = x * y, x and y read by their strides, out contiguous."""
+    pid = tl.program_id(0).to(tl.int64)
+    offsets = 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 = offsets < n_elements
+    remaining = offsets
+    x_off = tl.zeros_like(offsets)
+    y_off = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        xs = tl.load(x_stride_ptr + dim)
+        ys = tl.load(y_stride_ptr + dim)
+        idx = remaining % d
+        remaining = remaining // d
+        x_off = x_off + idx * xs
+        y_off = y_off + idx * ys
+    x = tl.load(x_ptr + x_off, mask=mask)
+    y = tl.load(y_ptr + y_off, mask=mask)
+    tl.store(output_ptr + offsets, x * y, mask=mask)
+
+
+@triton.jit
+def sub_strided_nd_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    alpha,
+    shape_ptr,
+    x_stride_ptr,
+    y_stride_ptr,
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+):
+    """out = x - alpha * y, x and y read by their strides, out contiguous."""
+    pid = tl.program_id(0).to(tl.int64)
+    offsets = 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 = offsets < n_elements
+    remaining = offsets
+    x_off = tl.zeros_like(offsets)
+    y_off = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        xs = tl.load(x_stride_ptr + dim)
+        ys = tl.load(y_stride_ptr + dim)
+        idx = remaining % d
+        remaining = remaining // d
+        x_off = x_off + idx * xs
+        y_off = y_off + idx * ys
+    x = tl.load(x_ptr + x_off, mask=mask)
+    y = tl.load(y_ptr + y_off, mask=mask)
+    tl.store(output_ptr + offsets, x - alpha * y, mask=mask)
+
+
+@triton.jit
+def div_strided_nd_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    shape_ptr,
+    x_stride_ptr,
+    y_stride_ptr,
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+):
+    """out = x / y, x and y read by their strides, out contiguous."""
+    pid = tl.program_id(0).to(tl.int64)
+    offsets = 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 = offsets < n_elements
+    remaining = offsets
+    x_off = tl.zeros_like(offsets)
+    y_off = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        xs = tl.load(x_stride_ptr + dim)
+        ys = tl.load(y_stride_ptr + dim)
+        idx = remaining % d
+        remaining = remaining // d
+        x_off = x_off + idx * xs
+        y_off = y_off + idx * ys
+    x = tl.load(x_ptr + x_off, mask=mask)
+    y = tl.load(y_ptr + y_off, mask=mask)
+    tl.store(output_ptr + offsets, x / y, mask=mask)
+
+
+@triton.jit
+def add_scalar_strided_nd_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    shape_ptr,
+    x_stride_ptr,
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+):
+    """out = x + scalar, x read by its strides, out contiguous (add_scalar_kernel's expression)."""
+    pid = tl.program_id(0).to(tl.int64)
+    offsets = 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 = offsets < n_elements
+    remaining = offsets
+    x_off = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        xs = tl.load(x_stride_ptr + dim)
+        idx = remaining % d
+        remaining = remaining // d
+        x_off = x_off + idx * xs
+    x = tl.load(x_ptr + x_off, mask=mask)
+    tl.store(output_ptr + offsets, x + scalar, mask=mask)
+
+
+@triton.jit
+def mul_scalar_strided_nd_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    shape_ptr,
+    x_stride_ptr,
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+):
+    """out = x * scalar, x read by its strides, out contiguous (mul_scalar_kernel's expression)."""
+    pid = tl.program_id(0).to(tl.int64)
+    offsets = 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 = offsets < n_elements
+    remaining = offsets
+    x_off = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        xs = tl.load(x_stride_ptr + dim)
+        idx = remaining % d
+        remaining = remaining // d
+        x_off = x_off + idx * xs
+    x = tl.load(x_ptr + x_off, mask=mask)
+    tl.store(output_ptr + offsets, x * scalar, mask=mask)

nbxtensor/ops/bitwise_and.py

@@ -0,0 +1,26 @@
+"""Bitwise AND kernel — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+Logic: result = x & y
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def bitwise_and_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)
+
+    result = x & y
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/bitwise_left_shift.py

@@ -0,0 +1,47 @@
+"""Bitwise left shift — pure @triton.jit kernel.
+
+Computes: output = x << y (element-wise).
+
+Both inputs must be integer tensors.
+
+Extracted from FlagGems reference (bitwise_left_shift.py).
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def bitwise_left_shift_kernel(
+    x_ptr,
+    y_ptr,
+    output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Bitwise left shift: out = x << y."""
+    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)
+
+    result = x << y
+    tl.store(output_ptr + offset, result, mask=mask)
+
+@triton.jit
+def bitwise_left_shift_scalar_kernel(
+    x_ptr,
+    output_ptr,
+    n_elements,
+    shift_amount,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Bitwise left shift by scalar: out = x << shift_amount."""
+    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)
+    result = x << shift_amount
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/bitwise_not.py

@@ -0,0 +1,24 @@
+"""Bitwise NOT kernel — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+Logic: result = ~x
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def bitwise_not_forward_kernel(
+    x_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)
+
+    result = ~x
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/bitwise_or.py

@@ -0,0 +1,26 @@
+"""Bitwise OR kernel — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+Logic: result = x | y
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def bitwise_or_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)
+
+    result = x | y
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/bitwise_right_shift.py

@@ -0,0 +1,47 @@
+"""Bitwise right shift — pure @triton.jit kernel.
+
+Computes: output = x >> y (element-wise, arithmetic shift).
+
+Both inputs must be integer tensors.
+
+Extracted from FlagGems reference (bitwise_right_shift.py).
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def bitwise_right_shift_kernel(
+    x_ptr,
+    y_ptr,
+    output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Bitwise right shift: out = x >> y."""
+    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)
+
+    result = x >> y
+    tl.store(output_ptr + offset, result, mask=mask)
+
+@triton.jit
+def bitwise_right_shift_scalar_kernel(
+    x_ptr,
+    output_ptr,
+    n_elements,
+    shift_amount,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Bitwise right shift by scalar: out = x >> shift_amount."""
+    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)
+    result = x >> shift_amount
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/bucketize.py

@@ -0,0 +1,37 @@
+"""Bucketize (searchsorted over a small boundary vector) — pure
+@triton.jit kernel.
+
+aten::bucketize semantics: for each element v, the returned index i
+counts the boundaries that sort strictly before v (right=False,
+``boundaries[i-1] < v <= boundaries[i]``) or before-or-equal
+(right=True, ``boundaries[i-1] <= v < boundaries[i]``).
+
+Boundary vectors on this path are decode-control scale (block-diagonal
+attention masks: one boundary per temporal frame / packed segment —
+tens of entries), so a linear scan per element beats a binary search
+and stays branch-free and bit-reproducible.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def bucketize_kernel(
+    x_ptr, boundaries_ptr, out_ptr,
+    n_elements, n_boundaries,
+    RIGHT: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    offs = 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)
+    in_range = offs < n_elements
+    x = tl.load(x_ptr + offs, mask=in_range, other=0)
+    acc = tl.zeros((BLOCK_SIZE,), dtype=tl.int64)
+    for i in range(0, n_boundaries):
+        b = tl.load(boundaries_ptr + i)
+        if RIGHT:
+            acc += (b <= x).to(tl.int64)
+        else:
+            acc += (b < x).to(tl.int64)
+    tl.store(out_ptr + offs, acc, mask=in_range)

nbxtensor/ops/cat_op.py

@@ -0,0 +1,87 @@
+"""Concatenation — pure @triton.jit kernel.
+
+Ported from FlagGems cat_copy_func_kernel_4 (Apache-2.0 license).
+Copies up to 4 input tensors per launch into the correct position
+in the output tensor along the cat dimension.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def cat_copy_kernel_4(
+    out_ptr,
+    in_ptr_a,
+    in_ptr_b,
+    in_ptr_c,
+    in_ptr_d,
+    dim_size_in_a,
+    dim_size_in_b,
+    dim_size_in_c,
+    dim_size_in_d,
+    dim_size_out,
+    dim_prod_post,
+    dim_offset_a,
+    dim_offset_b,
+    dim_offset_c,
+    dim_offset_d,
+    total_elements_a,
+    total_elements_b,
+    total_elements_c,
+    total_elements_d,
+    BLOCK_X: tl.constexpr,
+):
+    """Copy up to 4 input tensors into their correct positions in the output.
+
+    pid_x indexes blocks within a single input tensor.
+    pid_y selects which of the 4 input tensors this program handles.
+
+    The index remapping converts flat input indices to flat output indices
+    by splitting into (pre_dim, cat_dim, post_dim) components and applying
+    the cat dimension offset.
+    """
+    pid_x = tl.program_id(0).to(tl.int64)
+    pid_y = tl.program_id(1).to(tl.int64)
+
+    # Select the correct input tensor and its metadata
+    if pid_y == 0:
+        in_ptr = in_ptr_a
+    elif pid_y == 1:
+        in_ptr = in_ptr_b
+    elif pid_y == 2:
+        in_ptr = in_ptr_c
+    else:
+        in_ptr = in_ptr_d
+
+    # The per-input scalars are selected OUTSIDE the pointer branches, by tl.where,
+    # which unifies their types: Triton specialises a Python int as int32 when it
+    # fits, int64 when it does not, and as a constexpr when it equals 1 — so one
+    # input past 2^31 elements beside a small one made the if/elif branches
+    # disagree (`Mismatched type for total_elements`, Open-Sora-v2 VAE decode,
+    # 2026-09-25), and a `.to(tl.int64)` inside the branches broke on the
+    # constexpr form ('int' object has no attribute 'to', Flex.1-alpha).
+    dim_size_in = tl.where(pid_y == 0, dim_size_in_a, tl.where(pid_y == 1, dim_size_in_b, tl.where(pid_y == 2, dim_size_in_c, dim_size_in_d)))
+    dim_offset = tl.where(pid_y == 0, dim_offset_a, tl.where(pid_y == 1, dim_offset_b, tl.where(pid_y == 2, dim_offset_c, dim_offset_d)))
+    total_elements = tl.where(pid_y == 0, total_elements_a, tl.where(pid_y == 1, total_elements_b, tl.where(pid_y == 2, total_elements_c, total_elements_d)))
+
+    block_start = pid_x * BLOCK_X
+    offsets = tl.arange(0, BLOCK_X)
+    mask = block_start + offsets < total_elements
+
+    idx = block_start + offsets
+
+    # Decompose flat index into (pre, dim, post) coordinates
+    pre_idx = idx // (dim_size_in * dim_prod_post)
+    dim_idx = (idx // dim_prod_post) % dim_size_in
+    post_idx = idx % dim_prod_post
+
+    # Compute output index with the cat dimension offset applied
+    out_idx = (
+        pre_idx * dim_size_out * dim_prod_post
+        + (dim_idx + dim_offset) * dim_prod_post
+        + post_idx
+    )
+
+    data = tl.load(in_ptr + idx, mask=mask)
+    tl.store(out_ptr + out_idx, data, mask=mask)

nbxtensor/ops/celu.py

@@ -0,0 +1,20 @@
+"""CELU activation — pure @triton.jit kernel. Logic from FlagGems."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def celu_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    alpha,
+    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)
+    x_fp32 = x.to(tl.float32)
+    out = tl.where(x_fp32 > 0, x_fp32, alpha * (tl.exp(x_fp32 / alpha) - 1.0))
+    tl.store(output_ptr + offset, out, mask=mask)