nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/relu.py

@@ -0,0 +1,18 @@
+"""ReLU — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def relu_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)
+    out = tl.maximum(0, x)
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/remainder.py

@@ -0,0 +1,46 @@
+"""Remainder (modulo) — pure @triton.jit kernel.
+
+torch.remainder / Python `%` define the result to follow the sign of the
+DIVISOR: remainder(a, b) = a - floor(a / b) * b, so the result lies in [0, b)
+for b > 0 (e.g. remainder(-0.1, 1) = 0.9). Triton's `%` operator on floats is C
+`fmod` (trunc division), which follows the sign of the DIVIDEND (fmod(-0.1, 1) =
+-0.1). Using `%` silently broke any negative-input modulo — surfaced by the
+Kokoro SineGen phase wrap `(f0*harmonics/sr) % 1`, where the predictor's
+unconstrained (sometimes negative) F0 produced negative phases that never
+wrapped into [0,1), degrading the harmonic source. The floor formula matches
+torch for floats and for integers (the result is integer-valued, cast back
+cleanly); x/y promotes to float so `floor` is well-defined for int inputs too.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def remainder_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 - tl.math.floor((x / y).to(tl.float32)).to(x.dtype) * y
+    tl.store(output_ptr + offset, result, mask=mask)
+
+@triton.jit
+def remainder_scalar_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    divisor,
+    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.math.floor((x / divisor).to(tl.float32)).to(x.dtype) * divisor
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/repeat_interleave.py

@@ -0,0 +1,48 @@
+"""repeat_interleave — pure @triton.jit kernel.
+
+Handles the tensor-repeats case: given a 1D repeats tensor and its cumsum,
+writes indices so that element i is repeated repeats[i] times.
+
+Example: repeats=[2,3,1], cumsum=[2,5,6] -> output=[0,0,1,1,1,2]
+
+The self-int case (scalar repeats) is a stride trick handled in the wrapper.
+
+Ported from FlagGems repeat_interleave.py. Stripped FlagGems decorators.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def repeat_interleave_tensor_kernel(
+    repeats_ptr,
+    cumsum_ptr,
+    out_ptr,
+    size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Write source indices for tensor-based repeat_interleave.
+
+    Grid: (size,) — one program per input element.
+    Each program writes its index `pid` into repeats[pid] consecutive
+    positions starting at cumsum[pid] - repeats[pid].
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    mask = pid < size
+    cumsum = tl.load(cumsum_ptr + pid, mask, other=0)
+    # The loop bound below must be 32-bit: the Metal induction lowering does
+    # not terminate on a 64-bit `scf.for` bound (it refuses rather than hang).
+    # `repeats` arrives i64 by torch default. The cast is safe -- a per-element
+    # repeat count of 2**31 would need an 8 GB output for ONE element, which
+    # the allocator refuses long before this loop -- and numerically inert on
+    # CUDA: same iterations, narrower type. Addressing stays wide: `out_offset`
+    # promotes back through `cumsum`.
+    repeats = tl.load(repeats_ptr + pid, mask, other=0).to(tl.int32)
+    out_offset = cumsum - repeats
+
+    base_ptr = out_ptr + out_offset
+    for start_k in range(0, repeats, BLOCK_SIZE):
+        offsets_k = start_k + tl.arange(0, BLOCK_SIZE)
+        mask_k = offsets_k < repeats
+        tl.store(base_ptr + offsets_k, pid, mask=mask_k)

nbxtensor/ops/repeat_op.py

@@ -0,0 +1,132 @@
+"""Repeat (tile) — pure @triton.jit kernel.
+
+Copies input data into output with tiling via modular indexing.
+For each output element, the source index is computed as:
+    src[i] = (output_idx[i] % input_shape[i]) for each dimension.
+
+This is a rank-generic kernel that handles up to 6 dimensions.
+Grid: (ceil(num_output_elements / BLOCK_SIZE),)
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def repeat_1d_kernel(
+    in_ptr, out_ptr,
+    in_stride0,
+    out_stride0,
+    out_shape0,
+    in_shape0,
+    num_tasks,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Repeat kernel for 1D tensors."""
+    pid = tl.program_id(0).to(tl.int64)
+    tid = 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 = tid < num_tasks
+
+    i0 = tid
+
+    # Compute source index via modular indexing
+    in_offset = (i0 % in_shape0) * in_stride0
+    out_offset = i0 * out_stride0
+
+    val = tl.load(in_ptr + in_offset, mask=mask)
+    tl.store(out_ptr + out_offset, val, mask=mask)
+
+
+@triton.jit
+def repeat_2d_kernel(
+    in_ptr, out_ptr,
+    in_stride0, in_stride1,
+    out_stride0, out_stride1,
+    out_shape0, out_shape1,
+    in_shape0, in_shape1,
+    num_tasks,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Repeat kernel for 2D tensors."""
+    pid = tl.program_id(0).to(tl.int64)
+    tid = 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 = tid < num_tasks
+
+    # Decompose flat index into multi-index
+    i1 = tid % out_shape1
+    i0 = tid // out_shape1
+
+    # Source via modular indexing
+    in_offset = (i0 % in_shape0) * in_stride0 + (i1 % in_shape1) * in_stride1
+    out_offset = i0 * out_stride0 + i1 * out_stride1
+
+    val = tl.load(in_ptr + in_offset, mask=mask)
+    tl.store(out_ptr + out_offset, val, mask=mask)
+
+
+@triton.jit
+def repeat_3d_kernel(
+    in_ptr, out_ptr,
+    in_stride0, in_stride1, in_stride2,
+    out_stride0, out_stride1, out_stride2,
+    out_shape0, out_shape1, out_shape2,
+    in_shape0, in_shape1, in_shape2,
+    num_tasks,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Repeat kernel for 3D tensors."""
+    pid = tl.program_id(0).to(tl.int64)
+    tid = 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 = tid < num_tasks
+
+    i2 = tid % out_shape2
+    tmp = tid // out_shape2
+    i1 = tmp % out_shape1
+    i0 = tmp // out_shape1
+
+    in_offset = (
+        (i0 % in_shape0) * in_stride0
+        + (i1 % in_shape1) * in_stride1
+        + (i2 % in_shape2) * in_stride2
+    )
+    out_offset = i0 * out_stride0 + i1 * out_stride1 + i2 * out_stride2
+
+    val = tl.load(in_ptr + in_offset, mask=mask)
+    tl.store(out_ptr + out_offset, val, mask=mask)
+
+
+@triton.jit
+def repeat_4d_kernel(
+    in_ptr, out_ptr,
+    in_stride0, in_stride1, in_stride2, in_stride3,
+    out_stride0, out_stride1, out_stride2, out_stride3,
+    out_shape0, out_shape1, out_shape2, out_shape3,
+    in_shape0, in_shape1, in_shape2, in_shape3,
+    num_tasks,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Repeat kernel for 4D tensors."""
+    pid = tl.program_id(0).to(tl.int64)
+    tid = 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 = tid < num_tasks
+
+    i3 = tid % out_shape3
+    tmp = tid // out_shape3
+    i2 = tmp % out_shape2
+    tmp = tmp // out_shape2
+    i1 = tmp % out_shape1
+    i0 = tmp // out_shape1
+
+    in_offset = (
+        (i0 % in_shape0) * in_stride0
+        + (i1 % in_shape1) * in_stride1
+        + (i2 % in_shape2) * in_stride2
+        + (i3 % in_shape3) * in_stride3
+    )
+    out_offset = (
+        i0 * out_stride0 + i1 * out_stride1
+        + i2 * out_stride2 + i3 * out_stride3
+    )
+
+    val = tl.load(in_ptr + in_offset, mask=mask)
+    tl.store(out_ptr + out_offset, val, mask=mask)

nbxtensor/ops/residual_chain.py

@@ -0,0 +1,374 @@
+"""
+S5 — DC-AE residual chain band-streamed execution.
+
+Implements the runtime wrapper that absorbs a long residual chain
+(detected by `OpLevelTilingEngine._detect_residual_chains`) into a
+single band-streamed compute. Reduces per-band intermediate footprint
+from 3× full-resolution buffers to 1× band-sized buffer.
+
+The chain pattern (validated on Sana 4Kpx VAE DC-AE blocks):
+
+    T_base (fork output)
+      |
+      |--- (chain): conv -> silu -> conv -> permute -> rms_norm -> add -> permute
+      |
+      +--- merge_add: chain_output + T_base = output
+
+For each band of H rows:
+  1. Slice T_base [h_in_start:h_in_end] with halo for conv kernels
+  2. Run the chain ops on the band (PyTorch ATen)
+  3. Trim halo from chain output
+  4. Add T_base band → write to output[h_start:h_end]
+
+Live tensor footprint per band: T_base (full) + output (full, fill in
+place) + ~1 band intermediate (small) → fits 16 GiB on Sana 4Kpx VAE
+4096² peak.
+
+R34 conformant: pattern detection is structural (`fork → linear≥3 →
+merge`), no model-name lookup. R30 dual-branch: compiled path uses
+PyTorch ATen here; triton path port will add NBX equivalent.
+
+S5 P-PRISM-NEVER-REFUSE v2 2026-05-13.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+
+
+class ChainSentinel:
+    """Zero-allocation sentinel returned by chain intermediate
+    interceptors. Carries no real data — the merge interceptor
+    reads the precomputed merge output from the chain registry.
+    """
+
+    __slots__ = ("chain_id", "shape", "dtype", "device")
+
+    def __init__(self, chain_id, shape, dtype, device):
+        self.chain_id = chain_id
+        self.shape = tuple(shape)
+        self.dtype = dtype
+        self.device = device
+
+    @property
+    def ndim(self):
+        return len(self.shape)
+
+    @property
+    def _nbytes(self):
+        return 0
+
+    def __repr__(self):
+        return f"ChainSentinel(chain_id={self.chain_id}, shape={self.shape})"
+
+
+def resolve_chain_weights(
+    chain_spec: Dict[str, Any],
+    dag: Dict[str, Any],
+    weights: Dict[str, Any],
+) -> Dict[str, Any]:
+    """Look up the weight tensors each chain op needs.
+
+    Parses the chain ops' `input_tensor_ids` from the DAG, identifies
+    `param::*` references, and pulls the matching entries from the
+    executor's `_weights` dict.
+
+    Returns a dict shaped for `band_streamed_chain_torch`:
+        {
+          "conv1_weight": tensor, "conv1_bias": tensor or None,
+          "conv1_padding": [1, 1], "conv1_stride": [1, 1],
+          "conv1_dilation": [1, 1], "conv1_groups": 1,
+          "conv2_weight": tensor, "conv2_bias": tensor or None,
+          "conv2_padding": ..., (same fields)
+          "norm_weight": tensor, "norm_eps": 1e-6,
+          "post_norm_bias": tensor or None,
+          "permute_forward": [0, 2, 3, 1],
+          "permute_backward": [0, 3, 1, 2],
+        }
+
+    Returns None if any required weight cannot be resolved.
+    """
+    ops = dag.get("ops", {})
+    chain_uids = chain_spec["chain_uids"]
+    out: Dict[str, Any] = {
+        "conv1_weight": None, "conv1_bias": None,
+        "conv1_padding": [1, 1], "conv1_stride": [1, 1],
+        "conv1_dilation": [1, 1], "conv1_groups": 1,
+        "conv2_weight": None, "conv2_bias": None,
+        "conv2_padding": [1, 1], "conv2_stride": [1, 1],
+        "conv2_dilation": [1, 1], "conv2_groups": 1,
+        "norm_weight": None, "norm_eps": 1e-6,
+        "post_norm_bias": None,
+        "permute_forward": [0, 2, 3, 1],
+        "permute_backward": [0, 3, 1, 2],
+    }
+
+    conv_idx = 0
+    for uid in chain_uids:
+        op = ops.get(uid, {})
+        op_type = op.get("op_type", "")
+        in_tids = op.get("input_tensor_ids", [])
+        if op_type == "aten::convolution":
+            # input_tensor_ids: [activation, weight, optional bias]
+            if len(in_tids) >= 2:
+                w_tid = in_tids[1]
+                if isinstance(w_tid, str) and w_tid.startswith("param::"):
+                    w = weights.get(w_tid[len("param::"):])
+                    if w is None:
+                        return {}
+                    out[f"conv{conv_idx + 1}_weight"] = w
+            if len(in_tids) >= 3:
+                b_tid = in_tids[2]
+                if isinstance(b_tid, str) and b_tid.startswith("param::"):
+                    out[f"conv{conv_idx + 1}_bias"] = weights.get(
+                        b_tid[len("param::"):])
+            # Conv attrs from args: [act, w, b, stride, padding, dil,
+            # transposed, output_padding, groups]
+            attrs = op.get("attributes", {})
+            args = attrs.get("args", [])
+
+            def _get_list(arg):
+                if isinstance(arg, dict) and arg.get("type") == "list":
+                    return [int(v) if isinstance(v, int) else int(
+                        v.get("value", 1)) for v in arg.get("value", [])]
+                return None
+
+            def _get_scalar(arg):
+                if isinstance(arg, dict) and arg.get("type") == "scalar":
+                    return arg.get("value")
+                return None
+
+            if len(args) >= 4:
+                v = _get_list(args[3])
+                if v:
+                    out[f"conv{conv_idx + 1}_stride"] = v
+            if len(args) >= 5:
+                v = _get_list(args[4])
+                if v:
+                    out[f"conv{conv_idx + 1}_padding"] = v
+            if len(args) >= 6:
+                v = _get_list(args[5])
+                if v:
+                    out[f"conv{conv_idx + 1}_dilation"] = v
+            if len(args) >= 9:
+                v = _get_scalar(args[8])
+                if isinstance(v, int):
+                    out[f"conv{conv_idx + 1}_groups"] = v
+            conv_idx += 1
+        elif op_type == "custom::rms_norm":
+            if len(in_tids) >= 2:
+                w_tid = in_tids[1]
+                if isinstance(w_tid, str) and w_tid.startswith("param::"):
+                    w = weights.get(w_tid[len("param::"):])
+                    if w is None:
+                        return {}
+                    out["norm_weight"] = w
+            attrs = op.get("attributes", {})
+            for arg in attrs.get("args", []):
+                if isinstance(arg, dict) and arg.get("type") == "scalar":
+                    v = arg.get("value")
+                    if isinstance(v, float):
+                        out["norm_eps"] = v
+                        break
+        elif op_type == "aten::add":
+            # Post-rms_norm bias-add: a learnable shift.
+            # input_tensor_ids: [activation, bias_param]
+            if len(in_tids) >= 2:
+                b_tid = in_tids[1]
+                if isinstance(b_tid, str) and b_tid.startswith("param::"):
+                    out["post_norm_bias"] = weights.get(
+                        b_tid[len("param::"):])
+        elif op_type == "aten::permute":
+            attrs = op.get("attributes", {})
+            args = attrs.get("args", [])
+            if len(args) >= 2 and isinstance(args[1], dict) \
+                    and args[1].get("type") == "list":
+                perm = [int(v.get("value", 0)) if isinstance(v, dict)
+                        else int(v) for v in args[1].get("value", [])]
+                if perm == [0, 2, 3, 1]:
+                    out["permute_forward"] = perm
+                elif perm == [0, 3, 1, 2]:
+                    out["permute_backward"] = perm
+
+    if out["conv1_weight"] is None or out["conv2_weight"] is None:
+        return {}
+    return out
+
+
+
+# ---------------------------------------------------------------------------
+# NBX-PURE TRITON variant — R33 zero-torch port of band_streamed_chain_torch.
+# ---------------------------------------------------------------------------
+# Same band-streaming algorithm, same correctness invariant (halo_carry),
+# same memory shape. The only difference is the backend: every ATen op
+# is replaced by its NBX wrapper from `nbxtensor.ops.wrappers`.
+#
+# P-TRITON-LIVE-WATERMARK-AUDIT 2026-05-14 L4: closes the R30 chain wrapper
+# gap left open by commit c9d2581 ("future band_streamed_chain_nbx would
+# close this gap"). On Sana 4Kpx 16g triton the L2 LIVE_DUMP showed two
+# 4 GiB tensors at (1, 128, 4096, 4096) co-resident at add::89 because
+# the chain wrapper was SKIP'd on triton — natural NBX dispatch keeps
+# all chain intermediates fully materialized. With this wrapper the
+# chain runs band-by-band on the triton arena tensors, peak transient
+# drops to 1/tile_factor of the chain intermediate.
+
+
+def band_streamed_chain_nbx(
+    t_base: "Any",
+    chain_weights: Dict[str, Any],
+    tile_factor: int,
+    halo: int,
+) -> "Any":
+    """NBXTensor band-streamed residual chain (R33-pure mirror of
+    `band_streamed_chain_torch`).
+
+    Args:
+      t_base: NBXTensor, the fork tensor (residual base). Mutated in
+        place to hold the merge output. Returned for caller convenience.
+      chain_weights: dict from `resolve_chain_weights` — same shape as
+        for the torch variant, but weight values are NBXTensor when
+        the executor is in triton mode.
+      tile_factor: number of bands along H.
+      halo: rows of halo on each side per band.
+
+    Returns: t_base (mutated).
+
+    R33: no `torch.*`, no `F.*`. All ops via
+    `nbxtensor.ops.wrappers` and NBXTensor methods.
+
+    Correctness invariant: halo_carry is a clone of the rows the NEXT
+    band needs as its top halo, captured BEFORE the current band's
+    `add_` overwrites them. The first band uses t_base directly because
+    no top halo has been overwritten yet.
+    """
+    from nbxtensor.tensor import NBXTensor
+    from nbxtensor.backend.runtime import set_device
+    from nbxtensor.ops.conv_wrappers import conv2d_wrapper
+    from nbxtensor.ops.elementwise_wrappers import silu as silu_w, add as add_w
+    from nbxtensor.ops.norm_wrappers import rms_norm as rms_norm_w
+
+    # P-TRITON-LIVE-WATERMARK-AUDIT L4b: explicit device sync at
+    # entry. P-TRITON-CHAIN-CPU-POINTER C1 2026-05-14: diagnostic
+    # NBX_CHAIN_DIAG showed 1024² and 2048² chains succeed but
+    # 4096² chains fail with "Pointer argument (at 0) cannot be
+    # accessed from Triton (cpu tensor?)" while t_base+weights all
+    # report dev=0 contig=True at chain entry. The failure must
+    # therefore happen at an intermediate allocation (slice +
+    # contiguous → 1 GiB band, conv2d output → 1 GiB) where the
+    # device-context drifts. C3 fix: call `set_device(t_base)`
+    # before EVERY conv2d_wrapper / rms_norm_w / add_w in the
+    # band loop — overhead is tiny (single cudaSetDevice per op)
+    # vs the chain crash.
+    set_device(t_base)
+
+    N, C, H, W = t_base.shape
+    from nbxtensor.ops.tiling_sizes import residual_chain_bands  # the one band geometry (the census derives by it)
+    bands = residual_chain_bands(H, tile_factor, halo)
+    band_h = (H + tile_factor - 1) // tile_factor
+    # P-TRITON-CHAIN-CPU-POINTER C3c 2026-05-14: log halo / tile_factor /
+    # band_h / H so we can confirm the band_h+halo math matches what the
+    # diag CHAIN_STEP shows (e.g. 1030 = 1024 + 6 implies halo=6 not 2).
+    import os as _os_cd3
+    if _os_cd3.environ.get("NBX_CHAIN_DIAG", "0") == "1":
+        print(f"  [CHAIN_GEOM H={H} band_h={band_h} halo={halo} "
+              f"tile_factor={tile_factor}]", flush=True)
+
+    halo_carry = None  # type: Optional[NBXTensor]
+
+    for i, (h_start, h_end, h_in_start, h_in_end) in enumerate(bands):
+
+        # Build the input band.
+        # If we have a halo_carry AND the top-halo region [h_in_start,
+        # h_start) has been overwritten by the previous band's add_,
+        # take the top from halo_carry and the rest from t_base.
+        if halo_carry is not None and h_in_start < h_start:
+            top_size = h_start - h_in_start
+            top_band = halo_carry[:, :, -top_size:, :]
+            rest_band = t_base[:, :, h_start:h_in_end, :]
+            band = NBXTensor.cat([top_band, rest_band], dim=2).contiguous()
+        else:
+            band = t_base[:, :, h_in_start:h_in_end, :].contiguous()
+
+        # Save halo for next band BEFORE overwriting (clone detaches
+        # from t_base storage).
+        if i + 1 < tile_factor:
+            halo_save_size = min(halo, h_end - h_start)
+            halo_carry = t_base[:, :,
+                                h_end - halo_save_size:h_end, :].clone()
+
+        # conv1 — anchor device per op (the band size at 4096² is
+        # 1 GiB; the allocator path inside NBXTensor.empty cudaSetDevice
+        # can drift Triton's active device).
+        set_device(band)
+        w1 = chain_weights["conv1_weight"]
+        b1 = chain_weights["conv1_bias"]
+        import os as _os_cd2
+        if _os_cd2.environ.get("NBX_CHAIN_DIAG", "0") == "1":
+            _bp = getattr(band, "_data_ptr", "?")
+            _bd = getattr(band, "_device_idx", "?")
+            _bc = band.is_contiguous() if hasattr(
+                band, 'is_contiguous') else "?"
+            print(f"  [CHAIN_STEP i={i} h_start={h_start} h_end={h_end} "
+                  f"h_in_start={h_in_start} h_in_end={h_in_end} "
+                  f"conv1_pre band ptr={_bp} dev={_bd} "
+                  f"contig={_bc} shape={tuple(band.shape)}]",
+                  flush=True)
+        band = conv2d_wrapper(
+            band, w1, b1,
+            stride=tuple(chain_weights["conv1_stride"]),
+            padding=tuple(chain_weights["conv1_padding"]),
+            dilation=tuple(chain_weights["conv1_dilation"]),
+            groups=chain_weights["conv1_groups"],
+        )
+
+        # silu
+        set_device(band)
+        band = silu_w(band)
+
+        # conv2
+        set_device(band)
+        w2 = chain_weights["conv2_weight"]
+        b2 = chain_weights["conv2_bias"]
+        band = conv2d_wrapper(
+            band, w2, b2,
+            stride=tuple(chain_weights["conv2_stride"]),
+            padding=tuple(chain_weights["conv2_padding"]),
+            dilation=tuple(chain_weights["conv2_dilation"]),
+            groups=chain_weights["conv2_groups"],
+        )
+
+        # permute NCHW → NHWC (R33-pure NBXTensor.permute)
+        band = band.permute(*chain_weights["permute_forward"]).contiguous()
+
+        # rms_norm over last (feature) dim
+        set_device(band)
+        norm_w = chain_weights["norm_weight"]
+        eps = float(chain_weights.get("norm_eps", 1e-6))
+        band = rms_norm_w(band, norm_w, eps)
+
+        # post-norm bias shift (if any)
+        post_bias = chain_weights.get("post_norm_bias")
+        if post_bias is not None:
+            set_device(band)
+            band = add_w(band, post_bias)
+
+        # permute NHWC → NCHW
+        band = band.permute(*chain_weights["permute_backward"]).contiguous()
+
+        # Trim halo from chain output, then add to t_base in place via
+        # __setitem__ + add. NBX add_inplace_nbx requires identical
+        # shapes AND contiguous target, but t_base[:, :, h_start:h_end, :]
+        # is a non-contig view (slicing on dim 2) — so we use the
+        # `add_w` + setitem pattern: materialize the band-sized sum,
+        # write it back. The setitem path uses _strided_scatter for
+        # the non-contig dst and a single contiguous src copy.
+        trim_top = h_start - h_in_start
+        trim_h = h_end - h_start
+        band_out = band[:, :, trim_top:trim_top + trim_h, :]
+        t_base_slice = t_base[:, :, h_start:h_end, :]
+        # add returns a new contiguous tensor; setitem writes back.
+        set_device(t_base)
+        t_base[:, :, h_start:h_end, :] = add_w(t_base_slice, band_out)
+
+    return t_base
+

nbxtensor/ops/rmsnorm.py

@@ -0,0 +1,68 @@
+"""RMSNorm — 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.rms_norm_launches`)
+@nbx_autotune(configs=row_program_configs(), key=['BATCH_BUCKET', 'FEAT_BUCKET', 'scale_by_weight', 'SATURATE_F16'],
+              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 rms_norm_forward_kernel(
+    input_ptr, weight_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,
+    BLOCK_SIZE_BATCH: tl.constexpr,
+    BLOCK_SIZE_FEAT: tl.constexpr,
+    SATURATE_F16: tl.constexpr = False,   # the protected fp16 store: finite values clamp to ±65504
+):
+    """RMS normalization forward.
+
+    input: [batch_dim, feat_dim]
+    weight: [feat_dim] (optional)
+    output: [batch_dim, feat_dim]
+    """
+    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)
+    inv_rms = tl.rsqrt(tl.sum(inp * inp, axis=1) / feat_dim + eps)
+    output = inp * inv_rms[:, None]
+
+    if scale_by_weight:
+        weight = tl.load(weight_ptr + feat_offset, mask=feat_mask)
+        output *= weight
+
+    if SATURATE_F16:
+        # The store narrows to fp16: the same protected conversion the dtype
+        # engine's cast copy applied when this kernel stored fp32 and the wrap
+        # cast the result (copy_kernel's expression) — an fp16 stream with
+        # outliers past 65504 (VibeVoice's Qwen2, 2026-09-08) clamps instead
+        # of overflowing to inf.
+        is_finite = (output == output) & (output != float("inf")) & (output != float("-inf"))
+        clamped = tl.minimum(tl.maximum(output, -65504.0), 65504.0)
+        output = tl.where(is_finite, clamped, output)
+    tl.store(output_ptr, output,
+             mask=batch_mask[:, None] & feat_mask[None, :])