nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/trace_op.py

@@ -0,0 +1,64 @@
+"""Trace — pure @triton.jit kernel.
+
+Sum of diagonal elements of a 2D matrix.
+Single program iterates over the diagonal in blocks, accumulating in
+float32 (or int64 for integer types) for numerical stability.
+
+Extracted from FlagGems (Apache 2.0).
+
+Grid: (1,)
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def trace_kernel(
+    inp_ptr,
+    out_ptr,
+    num_diag,
+    stride0,
+    stride1,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Sum diagonal elements of a 2D matrix.
+
+    Iterates over the diagonal in blocks of BLOCK_SIZE, accumulating
+    partial sums. The diagonal stride is stride0 + stride1 (step to
+    move from (i,i) to (i+1,i+1)).
+
+    Args:
+        inp_ptr: input 2D matrix pointer
+        out_ptr: scalar output pointer
+        num_diag: min(M, N) — number of diagonal elements
+        stride0: row stride of input
+        stride1: column stride of input
+    """
+    # Select accumulation dtype based on input type
+    inp_dtype = inp_ptr.type.element_ty
+    if inp_dtype.is_int():
+        acc_dtype = tl.int64
+        other_val = 0
+    elif inp_dtype == tl.float64:
+        acc_dtype = tl.float64
+        other_val = 0.0
+    else:
+        acc_dtype = tl.float32
+        other_val = 0.0
+
+    acc = tl.zeros((BLOCK_SIZE,), dtype=acc_dtype)
+
+    diag_stride = stride0 + stride1
+
+    for i in range(0, tl.cdiv(num_diag, BLOCK_SIZE)):
+        block_start = i * BLOCK_SIZE
+        current_indices = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = current_indices < num_diag
+
+        ptr_offsets = current_indices * diag_stride
+        vals = tl.load(inp_ptr + ptr_offsets, mask=mask, other=other_val)
+        acc += vals.to(acc_dtype)
+
+    final_sum = tl.sum(acc, axis=0)
+    tl.store(out_ptr, final_sum.to(out_ptr.type.element_ty))

nbxtensor/ops/tril.py

@@ -0,0 +1,64 @@
+"""Lower triangular — pure @triton.jit kernels. Adapted from FlagGems triu.
+
+tril_kernel: 2D version, iterates over N blocks per row.
+tril_batch_kernel: batched version, handles N-D tensors flattened to (batch, M*N).
+
+The only difference from triu: condition is >= (keep below diagonal) vs <= (keep above).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit(do_not_specialize=["diagonal"])
+def tril_kernel(
+    X,
+    Y,
+    M,
+    N,
+    diagonal,
+    M_BLOCK_SIZE: tl.constexpr,
+    N_BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    row = pid * M_BLOCK_SIZE + tl.arange(0, M_BLOCK_SIZE)[:, None]
+    m_mask = row < M
+    X += row * N
+    Y += row * N
+
+    for n_offset in range(0, N, N_BLOCK_SIZE):
+        cols = n_offset + tl.arange(0, N_BLOCK_SIZE)[None, :]
+        n_mask = cols < N
+        mask = m_mask & n_mask
+
+        x = tl.load(X + cols, mask, other=0.0)
+        y = tl.where(row + diagonal >= cols, x, 0.0)
+        tl.store(Y + cols, y, mask=mask)
+
+
+@triton.jit(do_not_specialize=["diagonal"])
+def tril_batch_kernel(
+    X,
+    Y,
+    batch,
+    MN,
+    N,
+    diagonal,
+    BATCH_BLOCK_SIZE: tl.constexpr,
+    MN_BLOCK_SIZE: tl.constexpr,
+):
+    batch_id = tl.program_id(0).to(tl.int64)
+    mn_id = tl.program_id(1).to(tl.int64)
+    row = batch_id * BATCH_BLOCK_SIZE + tl.arange(0, BATCH_BLOCK_SIZE)[:, None]
+    batch_mask = row < batch
+    X += row * MN
+    Y += row * MN
+
+    cols = mn_id * MN_BLOCK_SIZE + tl.arange(0, MN_BLOCK_SIZE)[None, :]
+    mn_mask = cols < MN
+    mask = batch_mask & mn_mask
+    x = tl.load(X + cols, mask, other=0.0)
+    m = cols // N
+    n = cols % N
+    y = tl.where(m + diagonal >= n, x, 0.0)
+    tl.store(Y + cols, y, mask=mask)

nbxtensor/ops/triu.py

@@ -0,0 +1,62 @@
+"""Upper triangular — pure @triton.jit kernels. Extracted from FlagGems.
+
+triu_kernel: 2D version, iterates over N blocks per row.
+triu_batch_kernel: batched version, handles N-D tensors flattened to (batch, M*N).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit(do_not_specialize=["diagonal"])
+def triu_kernel(
+    X,
+    Y,
+    M,
+    N,
+    diagonal,
+    M_BLOCK_SIZE: tl.constexpr,
+    N_BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0).to(tl.int64)
+    row = pid * M_BLOCK_SIZE + tl.arange(0, M_BLOCK_SIZE)[:, None]
+    m_mask = row < M
+    X += row * N
+    Y += row * N
+
+    for n_offset in range(0, N, N_BLOCK_SIZE):
+        cols = n_offset + tl.arange(0, N_BLOCK_SIZE)[None, :]
+        n_mask = cols < N
+        mask = m_mask & n_mask
+
+        x = tl.load(X + cols, mask, other=0.0)
+        y = tl.where(row + diagonal <= cols, x, 0.0)
+        tl.store(Y + cols, y, mask=mask)
+
+
+@triton.jit(do_not_specialize=["diagonal"])
+def triu_batch_kernel(
+    X,
+    Y,
+    batch,
+    MN,
+    N,
+    diagonal,
+    BATCH_BLOCK_SIZE: tl.constexpr,
+    MN_BLOCK_SIZE: tl.constexpr,
+):
+    batch_id = tl.program_id(0).to(tl.int64)
+    mn_id = tl.program_id(1).to(tl.int64)
+    row = batch_id * BATCH_BLOCK_SIZE + tl.arange(0, BATCH_BLOCK_SIZE)[:, None]
+    batch_mask = row < batch
+    X += row * MN
+    Y += row * MN
+
+    cols = mn_id * MN_BLOCK_SIZE + tl.arange(0, MN_BLOCK_SIZE)[None, :]
+    mn_mask = cols < MN
+    mask = batch_mask & mn_mask
+    x = tl.load(X + cols, mask, other=0.0)
+    m = cols // N
+    n = cols % N
+    y = tl.where(m + diagonal <= n, x, 0.0)
+    tl.store(Y + cols, y, mask=mask)

nbxtensor/ops/unfold_op.py

@@ -0,0 +1,37 @@
+"""unfold_backward (overlap-add) — pure @triton.jit kernel.
+
+aten::unfold_backward is the gradient of `Tensor.unfold`: it scatter-adds the
+unfolded frames back into the original signal, so overlapping positions are
+summed. This is exactly the iSTFT overlap-add (frames [.., N_frames, size] with
+hop `step` → signal [.., L]).
+
+For output position i: out[i] = Σ_{f} grad[f, j] where j = i - f*step and
+0 <= j < size. The covering frames are f = i//step - fo for a small fo range
+(~size/step frames overlap a given position).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def unfold_backward_1d_kernel(
+    grad_ptr,        # (B, N_frames, size) contiguous
+    out_ptr,         # (B, L) contiguous
+    N_frames, size, step, L,
+    MAX_OVERLAP: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid_b = tl.program_id(0).to(tl.int64)
+    pid_l = tl.program_id(1).to(tl.int64)
+    offs = pid_l.to(tl.int64) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)   # 64-bit from the program id (register 58)
+    mask = offs < L
+    acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    base = pid_b * (N_frames * size)
+    for fo in range(MAX_OVERLAP):
+        f = (offs // step) - fo
+        j = offs - f * step
+        valid = (f >= 0) & (f < N_frames) & (j >= 0) & (j < size) & mask
+        val = tl.load(grad_ptr + base + f * size + j, mask=valid, other=0.0)
+        acc += tl.where(valid, val, 0.0)
+    tl.store(out_ptr + pid_b * L + offs, acc, mask=mask)

nbxtensor/ops/unique_op.py

@@ -0,0 +1,97 @@
+"""Unique — pure @triton.jit kernels.
+
+NOTE: The full unique implementation (FlagGems) requires a multi-phase
+pipeline (sort -> local-ne -> cumsum -> scatter -> fan-out) orchestrated
+by Python wrapper code that allocates intermediate buffers and chains
+multiple kernel launches. The core algorithmic kernels are extracted below.
+
+These kernels assume the input is PRE-SORTED. The Python wrapper must:
+1. Sort the input: sorted_data, sorted_indices = sort(input.ravel())
+2. Launch simple_unique_flat_kernel (for small inputs) or the
+   local_ne -> global_cumsum pipeline (for large inputs)
+3. Slice the output to [:unique_size]
+
+Extracted from FlagGems (Apache 2.0).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def simple_unique_flat_kernel(
+    sorted_data_ptr,
+    sorted_indices_ptr,
+    data_out_ptr,
+    inverse_indices_ptr,
+    idx_ptr,
+    unique_size_ptr,
+    return_inverse: tl.constexpr,
+    return_counts: tl.constexpr,
+    num_tasks,
+    tile_size: tl.constexpr,
+):
+    """Single-block unique on pre-sorted data (for small inputs <= 8192).
+
+    Compares each element with its predecessor to detect boundaries,
+    uses cumsum for compacted output indices, then scatters unique values.
+
+    Grid: (1,)
+    """
+    i0 = tl.arange(0, tile_size)
+    mask = i0 < num_tasks
+
+    # Load current and previous elements
+    a = tl.load(sorted_data_ptr + i0, mask=mask)
+    i0_prev = tl.where(i0 > 0, i0 - 1, 0)
+    b = tl.load(sorted_data_ptr + i0_prev, mask=mask)
+
+    # Detect unique boundaries: ne(a, prev(a)), first element always unique
+    ne_result = tl.where(i0 > 0, a != b, 0)
+    cumsum = tl.cumsum(ne_result)
+
+    # Write unique_size (last element's cumsum = num_unique - 1)
+    unique_size_mask = i0 == tile_size - 1
+    tl.store(unique_size_ptr + tl.zeros_like(i0), cumsum, mask=unique_size_mask)
+
+    # Scatter unique values to compacted output
+    tl.store(data_out_ptr + cumsum, a, mask=mask)
+
+    # Scatter inverse indices: original position -> compacted index
+    if return_inverse:
+        sorted_indices = tl.load(sorted_indices_ptr + i0, mask=mask)
+        tl.store(inverse_indices_ptr + sorted_indices, cumsum, mask=mask)
+
+    # Record boundary positions for count computation
+    if return_counts:
+        idx_mask = ((i0 == 0) | ne_result.to(tl.int1)) & mask
+        tl.store(idx_ptr + cumsum, i0, mask=idx_mask)
+
+
+@triton.jit
+def output_counts_kernel(
+    idx_ptr,
+    origin_num_tasks,
+    counts_ptr,
+    num_tasks,
+    tile_size: tl.constexpr,
+):
+    """Compute counts from boundary indices.
+
+    counts[i] = idx[i+1] - idx[i], with the last count going to the end.
+
+    Grid: (ceil(num_unique / tile_size),)
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    r = tl.arange(0, tile_size)
+    i0 = pid * tile_size + r
+    mask = i0 < num_tasks
+
+    idx = tl.load(idx_ptr + i0, mask=mask)
+
+    i0_next = i0 + 1
+    next_mask = i0_next < num_tasks
+    idx_next = tl.load(idx_ptr + i0_next, mask=next_mask)
+
+    counts = tl.where(i0_next < num_tasks, idx_next - idx, origin_num_tasks - idx)
+    tl.store(counts_ptr + i0, counts, mask=mask)

nbxtensor/ops/upsample_bilinear2d.py

@@ -0,0 +1,106 @@
+"""Bilinear 2D upsampling — pure @triton.jit kernel.
+
+For each output pixel (n, c, oh, ow), compute source coordinates in
+input space, find the 4 nearest input pixels, and bilinearly interpolate.
+
+Supports align_corners=True (maps corner-to-corner) and align_corners=False
+(maps center-to-center with half-pixel offset).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def upsample_bilinear2d_kernel(
+    input_ptr,
+    output_ptr,
+    N, C,
+    IH, IW,
+    OH, OW,
+    # Precomputed scale factors: how input coords map to output coords
+    # For align_corners=False: scale_h = IH / OH, scale_w = IW / OW
+    # For align_corners=True:  scale_h = (IH-1)/(OH-1), scale_w = (IW-1)/(OW-1)
+    scale_h,
+    scale_w,
+    # 0.0 for align_corners=True, 0.5 for align_corners=False
+    offset,
+    BLOCK_X: tl.constexpr,
+    BLOCK_Y: tl.constexpr,
+):
+    """Bilinear interpolation upsampling.
+
+    Grid: (cdiv(OW, BLOCK_X), cdiv(OH, BLOCK_Y), N*C)
+    Each program handles a BLOCK_Y x BLOCK_X tile of output pixels for one (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]
+
+    # Map output coords to input coords (floating point)
+    # align_corners=False: src = (dst + 0.5) * scale - 0.5
+    # align_corners=True:  src = dst * scale
+    src_h = (oh.to(tl.float32) + offset) * scale_h - offset  # [BLOCK_Y]
+    src_w = (ow.to(tl.float32) + offset) * scale_w - offset  # [BLOCK_X]
+
+    # Clamp to valid range
+    src_h = tl.maximum(src_h, 0.0)
+    src_w = tl.maximum(src_w, 0.0)
+    # IH/IW may arrive as a Python int (Triton-specialized) or a tl scalar
+    # depending on the call's shape — `.to()` fails on the former, float() on the
+    # latter. `* 1.0` promotes to float robustly for both.
+    src_h = tl.minimum(src_h, (IH - 1) * 1.0)
+    src_w = tl.minimum(src_w, (IW - 1) * 1.0)
+
+    # Floor and ceil indices
+    h0 = src_h.to(tl.int32)                         # [BLOCK_Y]
+    w0 = src_w.to(tl.int32)                         # [BLOCK_X]
+    h1 = tl.minimum(h0 + 1, IH - 1)                # [BLOCK_Y]
+    w1 = tl.minimum(w0 + 1, IW - 1)                # [BLOCK_X]
+
+    # Fractional parts (bilinear weights)
+    fh = src_h - h0.to(tl.float32)                  # [BLOCK_Y]
+    fw = src_w - w0.to(tl.float32)                  # [BLOCK_X]
+
+    # Bilinear weights for 4 corners
+    w_h0 = 1.0 - fh   # [BLOCK_Y]
+    w_h1 = fh          # [BLOCK_Y]
+    w_w0 = 1.0 - fw   # [BLOCK_X]
+    w_w1 = fw          # [BLOCK_X]
+
+    # Base offset for this (n, c) plane
+    base = pid_nc * IH * IW
+
+    # Masks
+    mask_h = oh < OH  # [BLOCK_Y]
+    mask_w = ow < OW  # [BLOCK_X]
+    mask_2d = mask_h[:, None] & mask_w[None, :]  # [BLOCK_Y, BLOCK_X]
+
+    # Load 4 corner values
+    # top-left:     input[n, c, h0, w0]
+    # top-right:    input[n, c, h0, w1]
+    # bottom-left:  input[n, c, h1, w0]
+    # bottom-right: input[n, c, h1, w1]
+    off_tl = base + h0[:, None] * IW + w0[None, :]
+    off_tr = base + h0[:, None] * IW + w1[None, :]
+    off_bl = base + h1[:, None] * IW + w0[None, :]
+    off_br = base + h1[:, None] * IW + w1[None, :]
+
+    v_tl = tl.load(input_ptr + off_tl, mask=mask_2d, other=0.0).to(tl.float32)
+    v_tr = tl.load(input_ptr + off_tr, mask=mask_2d, other=0.0).to(tl.float32)
+    v_bl = tl.load(input_ptr + off_bl, mask=mask_2d, other=0.0).to(tl.float32)
+    v_br = tl.load(input_ptr + off_br, mask=mask_2d, other=0.0).to(tl.float32)
+
+    # Bilinear interpolation
+    # result = w_h0 * (w_w0 * v_tl + w_w1 * v_tr) + w_h1 * (w_w0 * v_bl + w_w1 * v_br)
+    top = w_w0[None, :] * v_tl + w_w1[None, :] * v_tr
+    bot = w_w0[None, :] * v_bl + w_w1[None, :] * v_br
+    result = w_h0[:, None] * top + w_h1[:, None] * bot
+
+    # Store output
+    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/upsample_nearest2d.py

@@ -0,0 +1,52 @@
+"""Upsample nearest 2D — pure @triton.jit kernel.
+
+Extracted from FlagGems (FlagOpen/FlagGems) upsample_nearest2d.py.
+Stripped FlagGems-specific decorators (@libentry, runtime heuristics).
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def upsample_nearest2d_kernel(
+    ptr_o,
+    ptr_i,
+    N,
+    C,
+    OH,
+    OW,
+    IH,
+    IW,
+    s_n,
+    s_c,
+    s_h,
+    s_w,
+    reciprocal_scale_h,
+    reciprocal_scale_w,
+    BLOCK_SIZE: tl.constexpr,
+):
+    NC = N * C
+    nc_stride = tl.num_programs(axis=1)
+    nc_iter = (tl.program_id(axis=1).to(tl.int64)).to(tl.int32)   # loop bounds are 32-bit counts: the Metal lowering refuses a 64-bit scf.for bound (addressing stays 64-bit through the program ids)
+    pid = tl.program_id(axis=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)
+    ow = idx % OW
+    oh = idx // OW % OH
+
+    ih = tl.minimum((oh * reciprocal_scale_h).to(tl.int32), IH - 1)
+    iw = tl.minimum((ow * reciprocal_scale_w).to(tl.int32), IW - 1)
+
+    # The input is read through its strides: a first-frame branch feeds this op `x[:, :, 0]` of a
+    # 5-D activation, a view whose channel stride spans every frame — flat indexing read another
+    # channel's frames there (CogVideoX / Open-Sora frame 0). The output is fresh and contiguous.
+    offset_o = (nc_iter * OH + oh) * OW + ow
+    in_hw = ih.to(tl.int64) * s_h + iw.to(tl.int64) * s_w
+    dst_index_stride = nc_stride * OH * OW
+    while nc_iter < NC:
+        n = nc_iter // C
+        c = nc_iter % C
+        data = tl.load(ptr_i + n.to(tl.int64) * s_n + c.to(tl.int64) * s_c + in_hw)
+        tl.store(ptr_o + offset_o, data)
+        ptr_o += dst_index_stride
+        nc_iter += nc_stride

nbxtensor/ops/upsample_nearest3d.py

@@ -0,0 +1,61 @@
+"""Upsample nearest 3D — pure @triton.jit kernel.
+
+3D nearest neighbor upsampling: [N,C,D,H,W] -> [N,C,D*s_d,H*s_h,W*s_w].
+Adapted from our upsample_nearest2d.py by adding the depth dimension.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def upsample_nearest3d_kernel(
+    ptr_o,
+    ptr_i,
+    N,
+    C,
+    OD,
+    OH,
+    OW,
+    ID,
+    IH,
+    IW,
+    reciprocal_scale_d,
+    reciprocal_scale_h,
+    reciprocal_scale_w,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Nearest-neighbor 3D upsample.
+
+    Grid: (cdiv(OD*OH*OW, BLOCK_SIZE), cdiv(N*C, 4))
+    Each program handles BLOCK_SIZE output pixels across one (n,c) slice,
+    iterating over (n,c) in strides of nc_stride.
+    """
+    NC = N * C
+    nc_stride = tl.num_programs(axis=1)
+    nc_iter = (tl.program_id(axis=1).to(tl.int64)).to(tl.int32)   # loop bounds are 32-bit counts: the Metal lowering refuses a 64-bit scf.for bound (addressing stays 64-bit through the program ids)
+    pid = tl.program_id(axis=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)
+
+    # Decompose flat index -> (od, oh, ow)
+    ow = idx % OW
+    oh = (idx // OW) % OH
+    od = (idx // (OW * OH)) % OD
+
+    # Map output coords to input coords (nearest neighbor)
+    id_ = tl.minimum((od * reciprocal_scale_d).to(tl.int32), ID - 1)
+    ih = tl.minimum((oh * reciprocal_scale_h).to(tl.int32), IH - 1)
+    iw = tl.minimum((ow * reciprocal_scale_w).to(tl.int32), IW - 1)
+
+    # Compute offsets for contiguous [N, C, D, H, W] layout
+    offset_o = ((nc_iter * OD + od) * OH + oh) * OW + ow
+    offset_i = ((nc_iter * ID + id_) * IH + ih) * IW + iw
+    src_index_stride = nc_stride * ID * IH * IW
+    dst_index_stride = nc_stride * OD * OH * OW
+
+    while nc_iter < NC:
+        data = tl.load(ptr_i + offset_i)
+        tl.store(ptr_o + offset_o, data)
+        ptr_i += src_index_stride
+        ptr_o += dst_index_stride
+        nc_iter += nc_stride

nbxtensor/ops/var.py

@@ -0,0 +1,136 @@
+"""Variance — pure @triton.jit kernels.
+
+Welford online algorithm for numerically-stable variance.
+Global reduction via two-pass and dim-specific Welford kernel.
+Extracted from FlagGems var_mean (Apache-2.0).
+"""
+
+import triton
+import triton.language as tl
+
+
+# ---------------------------------------------------------------------------
+# Welford combine function (used by tl.reduce)
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def welford_func(mean_x, count_x, M_x, mean_y, count_y, M_y):
+    """Combine two Welford partial aggregates."""
+    count = count_x + count_y
+    _count = tl.maximum(count, 1)
+    mc_x = mean_x * count_x
+    mc_y = mean_y * count_y
+    mean = (mc_x + mc_y) / _count
+    M = M_x + mc_x * mean_x + M_y + mc_y * mean_y - count * mean * mean
+    return mean, count, M
+
+
+# ---------------------------------------------------------------------------
+# Global reduction: two-pass (partial stats → combine)
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def var_kernel_1(
+    X,
+    Acc,
+    Average,
+    Count,
+    N,
+    BLOCK_N: tl.constexpr,
+):
+    """Pass 1: compute partial sum, sum-of-squares, and count per block."""
+    pid = tl.program_id(0).to(tl.int64)
+    offset = pid * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    X = X + offset
+    Acc = Acc + pid
+    Average = Average + pid
+    Count = Count + pid
+    mask = offset < N
+
+    x = tl.load(X, mask, other=0.0).to(tl.float32)
+
+    count = tl.sum(mask.to(tl.float32))
+    average = tl.sum(x) / count
+    acc = tl.sum(x * x) - count * average * average
+
+    tl.store(Average, average)
+    tl.store(Acc, acc)
+    tl.store(Count, count)
+
+
+@triton.jit(do_not_specialize=["correction"])
+def var_kernel_2(
+    Acc,
+    Average,
+    Count,
+    Var,
+    N,
+    correction,
+    BLOCK_NUM,
+    BLOCK_N: tl.constexpr,
+):
+    """Pass 2: Welford-combine all blocks and output variance."""
+    offset = tl.arange(0, BLOCK_N)
+    mask = offset < BLOCK_NUM
+    Acc = Acc + offset
+    Average = Average + offset
+    Count = Count + offset
+    acc = tl.load(Acc, mask, other=0.0).to(tl.float32)
+    average = tl.load(Average, mask, other=0.0).to(tl.float32)
+    count = tl.load(Count, mask, other=0.0).to(tl.float32)
+
+    _mean, _, nvar = tl.reduce(
+        (average, count, acc), axis=0, combine_fn=welford_func
+    )
+
+    var = nvar / (N - correction)
+    tl.store(Var, var)
+
+
+# ---------------------------------------------------------------------------
+# Dim-specific: Welford online variance along one axis
+# ---------------------------------------------------------------------------
+
+@triton.jit(do_not_specialize=["correction"])
+def var_welford_kernel(
+    X,
+    Var,
+    M,
+    N,
+    correction,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Welford variance along compressed inner dim N.
+
+    X: [M, N] (dim-compressed) → Var: [M]
+    """
+    pid = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)[:, None]
+    X = X + pid * N
+    Var = Var + pid
+    row_mask = pid < M
+
+    _mean = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    _acc = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    _count = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+    for off in range(0, N, BLOCK_N):
+        cols = off + tl.arange(0, BLOCK_N)[None, :]
+        col_mask = cols < N
+        mask = row_mask & col_mask
+
+        x = tl.load(X + cols, mask, other=0.0).to(tl.float32)
+
+        count = _count + mask
+        cnt = tl.maximum(count, 1)
+        cur_mean = (_mean * _count + x) / cnt
+        _acc += (x - cur_mean) * (x - _mean) * mask
+        _mean = cur_mean
+        _count = count
+
+    _mean_r, _, acc_r = tl.reduce(
+        (_mean, _count, _acc), axis=1, combine_fn=welford_func
+    )
+    var = acc_r / (N - correction)
+    var = var[:, None]
+    tl.store(Var, var, row_mask)