nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/std.py

@@ -0,0 +1,113 @@
+"""Standard deviation — pure @triton.jit kernels.
+
+Global reduction via map-reduce (sum + sum_sq) and dim-specific two-pass
+(mean, then variance). Extracted from FlagGems (Apache-2.0).
+"""
+
+import triton
+import triton.language as tl
+
+
+# ---------------------------------------------------------------------------
+# Global reduction: map-reduce
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def std_map_kernel(X, Tmp_sum, Tmp_sum_sq, N, BLOCK_N: tl.constexpr):
+    """Pass 1: compute partial sum and sum-of-squares per block."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+    mask = offset < N
+    x = tl.load(X + offset, mask=mask, other=0.0).to(tl.float32)
+    sum_val = tl.sum(x, axis=0)
+    sum_sq_val = tl.sum(x * x, axis=0)
+    tl.store(Tmp_sum + pid, sum_val)
+    tl.store(Tmp_sum_sq + pid, sum_sq_val)
+
+
+@triton.jit
+def std_reduce_kernel(
+    Tmp_sum, Tmp_sum_sq, Out, N, correction, BLOCK_NUM, BLOCK_SIZE: tl.constexpr
+):
+    """Pass 2: combine blocks and compute std = sqrt(var)."""
+    total_sum_acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    total_sum_sq_acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    for off in range(0, BLOCK_NUM, BLOCK_SIZE):
+        offset = off + tl.arange(0, BLOCK_SIZE)
+        mask = offset < BLOCK_NUM
+        tmp_sum_vals = tl.load(Tmp_sum + offset, mask=mask, other=0.0).to(tl.float32)
+        tmp_sum_sq_vals = tl.load(Tmp_sum_sq + offset, mask=mask, other=0.0).to(
+            tl.float32
+        )
+        total_sum_acc += tmp_sum_vals
+        total_sum_sq_acc += tmp_sum_sq_vals
+    total_sum = tl.sum(total_sum_acc, axis=0)
+    total_sum_sq = tl.sum(total_sum_sq_acc, axis=0)
+    mean = total_sum / N
+    var = (total_sum_sq / N) - (mean * mean)
+    var = var * N / tl.maximum(N - correction, 1.0)
+    safe_var = tl.maximum(var, 0.0)
+    std_dev = tl.sqrt(safe_var)
+    tl.store(Out, std_dev.to(Out.dtype.element_ty))
+
+
+# ---------------------------------------------------------------------------
+# Dim-specific: two-pass (mean, then variance)
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def std_dim_kernel(
+    X,
+    Out,
+    stride_x_row,
+    stride_x_col,
+    M,
+    N,
+    correction,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Compute std along compressed inner dim N.
+
+    X: [M, N] (dim-compressed, strided) → Out: [M]
+    Two-pass: first compute mean, then sum of squared deviations.
+    """
+    pid_group = tl.program_id(0).to(tl.int64)
+    start_row = pid_group * BLOCK_M
+    row_offsets = start_row + tl.arange(0, BLOCK_M)
+    row_mask = row_offsets < M
+
+    # Pass 1: compute mean
+    mean_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    x_row_ptrs = X + row_offsets[:, None] * stride_x_row
+
+    for off in range(0, N, BLOCK_N):
+        col_offsets = off + tl.arange(0, BLOCK_N)
+        col_mask = col_offsets < N
+        x_ptrs = x_row_ptrs + col_offsets[None, :] * stride_x_col
+        final_mask = row_mask[:, None] & col_mask[None, :]
+        x = tl.load(x_ptrs, mask=final_mask, other=0.0)
+        mean_acc += x.to(tl.float32)
+
+    mean = tl.sum(mean_acc, axis=1) / N
+
+    # Pass 2: compute variance
+    var_acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for off in range(0, N, BLOCK_N):
+        col_offsets = off + tl.arange(0, BLOCK_N)
+        col_mask = col_offsets < N
+        x_ptrs = x_row_ptrs + col_offsets[None, :] * stride_x_col
+        final_mask = row_mask[:, None] & col_mask[None, :]
+        x = tl.load(x_ptrs, mask=final_mask, other=0.0)
+        diff = x.to(tl.float32) - mean[:, None]
+        var_acc += tl.where(final_mask, diff * diff, 0.0)
+
+    var = tl.sum(var_acc, axis=1)
+
+    denom = N - correction
+    var = var / tl.maximum(denom, 1e-12)
+    safe_var = tl.maximum(var, 0.0)
+    std_dev = tl.sqrt(safe_var)
+
+    out_ptrs = Out + row_offsets
+    tl.store(out_ptrs, std_dev.to(Out.dtype.element_ty), mask=row_mask)

nbxtensor/ops/strided_copy.py

@@ -0,0 +1,200 @@
+"""Strided ↔ contiguous copy kernels — N-D scalable, Triton-JIT.
+
+Two operations:
+
+* `strided_copy_kernel`    — read from a non-contiguous (strided) source,
+                             write to a contiguous destination. Used by
+                             `NBXTensor.clone` / `NBXTensor.contiguous`
+                             whenever the source is a view (permute,
+                             transpose, narrow, expand, ...).
+
+* `strided_scatter_kernel` — read from a contiguous source, write to a
+                             non-contiguous destination. Used by
+                             `NBXTensor.__setitem__` (KV cache indexed
+                             writes into a narrow view).
+
+* `strided_copy_nd_kernel`  — read from a strided source, write to a
+                             strided destination, converting the dtype on
+                             the store with `copy_kernel`'s exact protected
+                             fp16 rounding. One launch where `__setitem__`,
+                             `copy_` and `.to()` used to take a contiguous
+                             copy of the source, a cast of that copy and a
+                             scatter of the cast (three launches and two
+                             transients per KV-cache write, 2026-09-07).
+
+Both kernels are parameterised by `NDIM: tl.constexpr`. Triton
+specialises and caches one compiled kernel per distinct ndim
+encountered at runtime: the per-dimension index decomposition loop is
+unrolled at compile time via `tl.static_range(NDIM)`.
+
+The shape and stride vectors are passed as GPU pointers to small int64
+arrays (NDIM elements each). The kernel loads them inside the JIT body,
+so callers can handle any tensor up to PyTorch's 25-D hard limit
+without recompiling the code. A 6-D PixArt patchify clone, an 8-D
+SANA-Video VAE reshape, and a hypothetical 10-12-D multi-view 4D-scene
+permute all route through the same kernel text.
+
+Why this matters: the previous version was hardcoded to 5 dims
+(`d0..d4`, `s0..s4`). The wrapper silently dropped higher dims via
+`[1] * (5 - ndim)` returning `[]` for ndim > 5. 6-D PixArt clones
+produced out-of-bounds reads past the source allocation — silently
+corrupted data (Alpha) or `cudaErrorIllegalAddress` (Sigma), and
+8-D SANA-Video VAE clones had been unreachable entirely.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def strided_copy_kernel(
+    src_ptr, dst_ptr,
+    n_elements,
+    shape_ptr,     # GPU pointer → NDIM int64 shape values (source)
+    stride_ptr,    # GPU pointer → NDIM int64 stride values (source)
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+    WIDE: tl.constexpr = True,
+):
+    """Copy strided src → contiguous dst. Dimensions are unrolled at
+    compile time via the NDIM constexpr; shape/strides are loaded from
+    GPU scratch buffers so the kernel text is identical for every
+    ndim Triton specialises.
+
+    Args:
+        src_ptr:     Source tensor data pointer (may be non-contiguous).
+        dst_ptr:     Destination tensor data pointer (contiguous).
+        n_elements:  Total elements to copy (= src._numel).
+        shape_ptr:   Pointer to NDIM int64 shape values (source shape).
+        stride_ptr:  Pointer to NDIM int64 stride values (source strides).
+        BLOCK_SIZE:  Elements per thread block (constexpr).
+        NDIM:        Rank of the tensors (constexpr; triggers
+                     specialisation + kernel cache entry).
+    """
+    if WIDE:
+        offsets = tl.program_id(0).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)
+    else:
+        offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # Decompose the flat offset into multi-dim indices, walking dims
+    # from the innermost (NDIM-1) outward. At each step, `remaining %
+    # d` gives the current-dim index and `remaining // d` feeds the
+    # next outer dim. The src_offsets accumulator multiplies each
+    # per-dim index by the source's stride for that dim — this is
+    # what lets the kernel read from a non-contiguous layout while
+    # writing contiguously.
+    remaining = offsets
+    src_offsets = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        s = tl.load(stride_ptr + dim)
+        if not WIDE:
+            d = d.to(tl.int32)
+            s = s.to(tl.int32)
+        idx = remaining % d
+        remaining = remaining // d
+        src_offsets = src_offsets + idx * s
+
+    vals = tl.load(src_ptr + src_offsets, mask=mask)
+    tl.store(dst_ptr + offsets, vals, mask=mask)
+
+
+@triton.jit
+def strided_scatter_kernel(
+    src_ptr, dst_ptr,
+    n_elements,
+    shape_ptr,     # GPU pointer → NDIM int64 shape values (destination)
+    stride_ptr,    # GPU pointer → NDIM int64 stride values (destination)
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+    WIDE: tl.constexpr = True,
+):
+    """Scatter contiguous src → strided dst. Inverse of
+    `strided_copy_kernel`. Same N-D specialisation strategy.
+
+    Args:
+        src_ptr:     Source tensor data pointer (contiguous).
+        dst_ptr:     Destination tensor data pointer (may be non-
+                     contiguous — e.g. a narrow view into a KV cache).
+        n_elements:  Total elements to scatter (= dst._numel).
+        shape_ptr:   Pointer to NDIM int64 shape values (dst shape).
+        stride_ptr:  Pointer to NDIM int64 stride values (dst strides).
+        BLOCK_SIZE:  Elements per thread block.
+        NDIM:        Rank of the tensors (constexpr).
+    """
+    if WIDE:
+        offsets = tl.program_id(0).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)
+    else:
+        offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    remaining = offsets
+    dst_offsets = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        s = tl.load(stride_ptr + dim)
+        if not WIDE:
+            d = d.to(tl.int32)
+            s = s.to(tl.int32)
+        idx = remaining % d
+        remaining = remaining // d
+        dst_offsets = dst_offsets + idx * s
+
+    vals = tl.load(src_ptr + offsets, mask=mask)
+    tl.store(dst_ptr + dst_offsets, vals, mask=mask)
+
+
+@triton.jit
+def strided_copy_nd_kernel(
+    src_ptr, dst_ptr,
+    n_elements,
+    shape_ptr,        # GPU pointer → NDIM int64 extents (shared by src and dst)
+    src_stride_ptr,   # GPU pointer → NDIM int64 source strides (0 on a broadcast dim)
+    dst_stride_ptr,   # GPU pointer → NDIM int64 destination strides
+    BLOCK_SIZE: tl.constexpr,
+    NDIM: tl.constexpr,
+    WIDE: tl.constexpr = True,
+    SATURATE_F16: tl.constexpr = False,
+):
+    """Copy strided src → strided dst, casting on the store.
+
+    The flat element index is decomposed once into the multi-index and
+    each side's offset accumulates its own strides, so any pair of
+    layouts of the same extents copies in one pass — a broadcast source
+    (stride 0) included. SATURATE_F16 is the protected float16 downcast
+    of `copy_kernel`, the same expression, so a cast through this kernel
+    is bit-identical to a cast through `copy_kernel` after `contiguous()`.
+    """
+    if WIDE:
+        offsets = tl.program_id(0).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)
+    else:
+        offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    remaining = offsets
+    src_offsets = tl.zeros_like(offsets)
+    dst_offsets = tl.zeros_like(offsets)
+    for i in tl.static_range(NDIM):
+        dim = NDIM - 1 - i
+        d = tl.load(shape_ptr + dim)
+        ss = tl.load(src_stride_ptr + dim)
+        ds = tl.load(dst_stride_ptr + dim)
+        if not WIDE:
+            d = d.to(tl.int32)
+            ss = ss.to(tl.int32)
+            ds = ds.to(tl.int32)
+        idx = remaining % d
+        remaining = remaining // d
+        src_offsets = src_offsets + idx * ss
+        dst_offsets = dst_offsets + idx * ds
+
+    vals = tl.load(src_ptr + src_offsets, mask=mask)
+    if SATURATE_F16:
+        f = vals.to(tl.float32)
+        is_finite = (f == f) & (f != float("inf")) & (f != float("-inf"))
+        clamped = tl.minimum(tl.maximum(f, -65504.0), 65504.0)
+        vals = tl.where(is_finite, clamped, f)
+    tl.store(dst_ptr + dst_offsets, vals, mask=mask)

nbxtensor/ops/sub.py

@@ -0,0 +1,35 @@
+"""Subtraction — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def sub_forward_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    alpha,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x - alpha * 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)
+    tl.store(output_ptr + offset, x - alpha * y, mask=mask)
+
+@triton.jit
+def rsub_forward_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = scalar - x (reverse subtraction)"""
+    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, scalar - x, mask=mask)

nbxtensor/ops/sum.py

@@ -0,0 +1,40 @@
+"""Sum reduction — pure @triton.jit kernel, the row walked in tiles of the certified width."""
+
+import triton
+import triton.language as tl
+
+from nbxtensor.ops._configs import nbx_autotune, row_reduce_configs, prune_row_reduce
+
+
+@nbx_autotune(configs=row_reduce_configs(), key=['BATCH_BUCKET', 'FEAT_BUCKET'],
+              prune_configs_by={'early_config_prune': prune_row_reduce}, cache_results=True)
+@triton.jit
+def sum_forward_kernel(
+    input_ptr, output_ptr,
+    batch_dim, feat_dim, BATCH_BUCKET, FEAT_BUCKET,
+    input_batch_stride, input_feat_stride,
+    BLOCK_SIZE_BATCH: tl.constexpr,
+    BLOCK_SIZE_FEAT: tl.constexpr,
+):
+    """Sum reduction over the last dimension: input [batch_dim, feat_dim] -> output [batch_dim].
+
+    The first tile is reduced alone and every further tile combined into it, so a row that fits one
+    tile is reduced by exactly the instructions it was before the loop existed.
+    """
+    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
+    row_ptr = input_ptr + input_batch_stride * batch_offset[:, None]
+
+    feat_mask = feat_offset < feat_dim
+    inp = tl.load(row_ptr + input_feat_stride * feat_offset[None, :],
+                  mask=batch_mask[:, None] & feat_mask[None, :], other=0.0).to(tl.float32)
+    result = tl.sum(inp, axis=1)
+    for start in range(BLOCK_SIZE_FEAT, feat_dim, BLOCK_SIZE_FEAT):
+        offs = start + feat_offset
+        feat_mask = offs < feat_dim
+        inp = tl.load(row_ptr + input_feat_stride * offs[None, :].to(tl.int64),
+                      mask=batch_mask[:, None] & feat_mask[None, :], other=0.0).to(tl.float32)
+        result += tl.sum(inp, axis=1)
+    tl.store(output_ptr + batch_offset, result, mask=batch_mask)

nbxtensor/ops/tan.py

@@ -0,0 +1,19 @@
+"""Tan (tangent) — pure @triton.jit kernel. Logic from FlagGems."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def tan_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)
+    x_fp32 = x.to(tl.float32)
+    out = tl.sin(x_fp32) / tl.cos(x_fp32)
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/tanh.py

@@ -0,0 +1,21 @@
+"""Tanh — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+from ._common import sigmoid
+
+@triton.jit
+def tanh_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)
+    x_fp32 = x.to(tl.float32)
+    out = 2.0 * sigmoid(2.0 * x_fp32) - 1.0
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/threshold.py

@@ -0,0 +1,19 @@
+"""Threshold activation — pure @triton.jit kernel. Logic from FlagGems."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def threshold_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    threshold_val, 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)
+    out = tl.where(x > threshold_val, x, value)
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/tiling.yml

@@ -0,0 +1,91 @@
+# The TilingEngine's sizes — every engine-wide value a memory split is cut by.
+#
+# Read by the TilingEngine's sizing half and nowhere else — `nbxtensor/ops/tiling_sizes.py`, for
+# both engine branches; a missing key is refused by name. Values that belong to an ARCHITECTURE stay in its hardware
+# profile (config/vendors/<vendor>/<arch>.yml): memory.sdpa_math_max_scores_bytes,
+# memory.sdpa_math_scores_device_fraction, memory.sdpa_math_min_chunk_rows,
+# memory.sdpa_math_max_chunks, tiling.extent_lattice.
+#
+# Each value carries its SOURCE: the reasoning, measurement or workstream that set it.
+
+conv3d:
+  # 1 GiB. The per-transient bound of the chunked temporal path, and the fold size above which
+  # the one-shot path's peak is priced at all: small and validated convolutions never chunk
+  # (P-TRITON-LIVE-SET-AUDIT; Allegro-TI2V vae_encoder at 720x1280x12f fp32, 5.66 GB a plane ->
+  # 30.4 GB live on a 32 GB V100 in triton only; the decision made Prism's, 2026-09-29).
+  chunk_bytes: 1073741824
+  # 256 MiB. The fixed slack added to the one-shot path's priced peak (conv3d_need).
+  slack_bytes: 268435456
+
+conv2d:
+  # 4 GiB. A convolution whose output exceeds it streams in bands along H; it leaves headroom for
+  # weights, activations and arena overhead on a V100 32 GB (P-SANA-4KPX-RUNTIME Step 1). A
+  # one-shot conv3d whose conv2d output exceeds it holds one band's worth more.
+  band_bytes: 4294967296
+  # Bands of about half the band budget: the other half covers the transient input slice and the
+  # kernel's intermediate (P-SANA-4KPX-RUNTIME Step 1, `_conv2d_band_streamed`).
+  band_budget_divisor: 2
+
+sdpa:
+  # 2 GiB. The fp32 scores bound a non-power-of-two head dim routes on when the architecture
+  # declares no `memory.sdpa_math_max_scores_bytes` (the scores budget LLM-class prefills chunk
+  # under a handful of times: 4k context -> 2 chunks, 16k -> 16; P-NONDET-LONG-ROW).
+  non_pow2_head_scores_bytes: 2147483648
+
+op_level:
+  # 0.85 of the card, minus the component's resident weights: the budget an op's transient is
+  # measured against — overflow detection, the op-level plan and the conv3d chunk decision share
+  # it, or they disagree about which upsamples are free (`_tiling_aware_zero_alloc_uids`).
+  op_budget_card_fraction: 0.85
+  # 0.65 of the card, minus the conv output that stays full-resolution allocated: one band's
+  # budget for an upsample->conv fusion pair or a standalone convolution (S5, P-SANA-4KPX).
+  conv_band_card_fraction: 0.65
+  # 256 MB: the smallest per-band budget a band factor is computed against (a budget under it,
+  # "< 256 MB usable", is raised to it).
+  band_budget_floor_bytes: 268435456
+  # A tiled op runs in at least two bands.
+  min_band_factor: 2
+  # Band factors round up to a power of two for halo alignment, at most 64.
+  max_band_factor: 64
+  # An upsample whose output exceeds 25 % of the card is worth fusing into its conv even when the
+  # conv does not overflow (typical: 16 GB on a 32 GB card).
+  upsample_overflow_card_fraction: 0.25
+  rms_norm:
+    # 20 % of the card: an rms_norm whose output exceeds it runs in H bands — a stricter bound
+    # skipped the exactly-8 GB tensors of Sana 4Kpx on a 32 GB card.
+    overflow_card_fraction: 0.20
+    # One band's budget: 0.10 of the card (tightened from 0.65 because S5 chains carry 5-7 GiB of
+    # state into the rms_norm dispatch and each band's fp32 cast must fit beside it; the comment
+    # beside the code said 0.20, the code has carried 0.10).
+    band_card_fraction: 0.10
+    # At least four bands.
+    min_band_factor: 4
+    # At most eight: sixteen produced a numerically wrong output at Sana 4Kpx's rms_norm::27
+    # (post-chain wrapper, NHWC view of non-contiguous storage); eight was validated on 32 GB.
+    max_band_factor: 8
+  residual_chain:
+    # A residual chain (fork -> >=3 ops -> merge) streams in bands of at most 1 GiB at runtime
+    # (P-PRISM-NEVER-REFUSE v2 S5, 2026-05-13).
+    band_target_bytes: 1073741824
+    # A chain streams in at least two bands (the solver's chain rule carried the conv band floor,
+    # 2, so a change of that floor moved this one silently; it is its own value now).
+    min_band_factor: 2
+    # Its band factor rounds up to a power of two, at most 32.
+    max_band_factor: 32
+    # 1.6 GB: a chain is detected only when its base tensor (fp32, conservative) reaches it —
+    # 1.6 GB at fp32, ~0.8 GB at bf16 (S5).
+    min_base_bytes_fp32: 1600000000
+
+overlap:
+  # A tile's overlap is its extent // 8, at least 4 to ensure some blending (TilingEngine's
+  # data-driven init; Prism's decode and encode tilings follow the same proportional law).
+  spatial_divisor: 8
+  spatial_min: 4
+  # A temporal tile overlaps its neighbour by t_tile // 4 frames (Prism's linear-class tiling).
+  temporal_divisor: 4
+
+inplace:
+  # 1 GiB. A residual add or an element-wise unary writes into the buffer of its dead input when
+  # its output reaches it (P-SANA-4KPX-RUNTIME 2026-05-05, multi-branch fusion fix vector B; the
+  # unary rewrite uses the same figure, real-esrgan-x8 at 1024x1024 on a 16 GB V100).
+  min_bytes: 1073741824

nbxtensor/ops/tiling_policy.py

@@ -0,0 +1,43 @@
+"""The tiling policy's default: the split sizes the tiling rules read (`tiling.yml`, beside this module).
+
+The library owns the file and reads it itself, so the rules run with no application around them. A caller may set
+its own table (`set_tiling`, plain data; `None` restores the library's); NeuroBrix reads this same file through the
+library's `tiling_sizes`. A missing file is refused by name: no split is ever cut by a size the data
+does not state.
+"""
+from __future__ import annotations
+
+import copy
+import functools
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+_FILE = Path(__file__).resolve().parent / "tiling.yml"
+
+__all__ = ["tiling_table", "set_tiling"]
+
+_SET: Dict[str, Any] = {}          # {"table": dict} once a caller set one
+
+
+@functools.lru_cache(maxsize=1)
+def _raw() -> Dict[str, Any]:
+    import yaml
+    if not _FILE.is_file():
+        raise FileNotFoundError(f"ZERO FALLBACK: the tiling policy is not found: {_FILE}")
+    return yaml.safe_load(_FILE.read_text()) or {}
+
+
+def tiling_table() -> Dict[str, Any]:
+    """The policy in force: the table a caller set, else the library's file (a copy: the caller may not edit it)."""
+    return copy.deepcopy(_SET["table"] if "table" in _SET else _raw())
+
+
+def set_tiling(table: Optional[Dict[str, Any]]) -> None:
+    """Set the split sizes the tiling rules read (a whole table, as `tiling.yml` writes it); `None` restores the
+    library's own."""
+    if table is None:
+        _SET.clear()
+    elif not isinstance(table, dict):
+        raise TypeError(f"set_tiling: a table is a dict as tiling.yml writes it, not {type(table).__name__}")
+    else:
+        _SET["table"] = copy.deepcopy(table)