nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/clamp.py

@@ -0,0 +1,24 @@
+"""Clamp — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def clamp_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    min_val, max_val,
+    has_min: tl.constexpr,
+    has_max: tl.constexpr,
+    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)
+    if has_min:
+        x = tl.maximum(min_val, x)
+    if has_max:
+        x = tl.minimum(max_val, x)
+    tl.store(output_ptr + offset, x, mask=mask)

nbxtensor/ops/clamp_min.py

@@ -0,0 +1,22 @@
+"""Clamp min — pure @triton.jit kernel.
+
+aten::clamp_min(x, min_val) = max(x, min_val). Separate file per ATen op convention.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def clamp_min_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    min_val,
+    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)
+    result = tl.maximum(min_val, x)
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/conv2d.py

@@ -0,0 +1,265 @@
+"""Conv2D forward — pure @triton.jit kernel with autotune.
+
+From attorch (BobMcDear/attorch) conv_kernels.py — chosen over FlagGems
+for V100 compatibility (proper num_stages handling). Im2col approach
+with tl.dot, fp32 accumulation, groups support.
+
+Phase 1.5 conv2d Step (b) (2026-05): @triton.autotune + cache_results=True
+adoption — same pattern as mm/bmm/addmm/baddbmm (commits d514bdb +
+20ed765). Configs restricted to Volta-viable subspace (BLOCK_BHW ∈
+{32,64,128}, BLOCK_OUTF ∈ {32,64,128}, BLOCK_INF ∈ {16,32,64}, num_warps
+∈ {2,4}, num_stages ∈ {2,3}). Triton compiler ejects configs that
+exceed sm_70's 96 KB SMEM at compile time (native fallback). See
+CLAUDE.md "Autotune policy" for the doctrinal exception scope.
+"""
+
+import triton
+
+from nbxtensor.ops._configs import nbx_autotune
+import triton.language as tl
+
+from nbxtensor.ops._autotune_policy import maybe_pin_single, is_conv2d_pinned
+
+
+# Volta-viable conv2d autotune subspace.
+# SMEM budget per config (rough): 2 stages × dtype_bytes × (BLOCK_BHW *
+# BLOCK_INF [input tile] + BLOCK_INF * BLOCK_OUTF [weight tile]).
+# fp16 (2B), 2 stages: BLOCK_BHW=128, BLOCK_INF=64, BLOCK_OUTF=128 →
+# 2 * 2 * (128*64 + 64*128) = 65 KB → fits 96 KB sm_70.
+# BLOCK_BHW=128, BLOCK_INF=64, BLOCK_OUTF=128, stages=3 → 98 KB → over.
+# Triton ejects over-SMEM configs at compile time via native fallback.
+_CONV2D_AUTOTUNE_VOLTA = [
+    triton.Config({'BLOCK_SIZE_BHW': 32,  'BLOCK_SIZE_OUTF': 32,  'BLOCK_SIZE_INF': 16}, num_stages=2, num_warps=2),
+    triton.Config({'BLOCK_SIZE_BHW': 32,  'BLOCK_SIZE_OUTF': 32,  'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=2),
+    triton.Config({'BLOCK_SIZE_BHW': 32,  'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=2),
+    triton.Config({'BLOCK_SIZE_BHW': 32,  'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 32}, num_stages=3, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 32,  'BLOCK_SIZE_OUTF': 128, 'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 32,  'BLOCK_SIZE_INF': 16}, num_stages=2, num_warps=2),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 32,  'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 32,  'BLOCK_SIZE_INF': 32}, num_stages=3, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 16}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 32}, num_stages=3, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 64}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 128, 'BLOCK_SIZE_INF': 16}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 64,  'BLOCK_SIZE_OUTF': 128, 'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 128, 'BLOCK_SIZE_OUTF': 32,  'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 128, 'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 16}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 128, 'BLOCK_SIZE_OUTF': 64,  'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_SIZE_BHW': 128, 'BLOCK_SIZE_OUTF': 128, 'BLOCK_SIZE_INF': 32}, num_stages=2, num_warps=4),
+]
+
+# Every tile in BOTH orientations (`SPATIAL_FAST`): the certifier picks per key per hardware. Metal takes the
+# spatial-fast tile (CogVideoX-2b VAE 3x3 (9,128,242x362) 109 -> 55 ms, Mac 2026-10-21); on V100 today's tile wins
+# five shapes of six (9-40 %, Dell 2026-10-21 01:53). Same key; an entry certified before the axis existed names no
+# orientation and keeps today's tile.
+_CONV2D_AUTOTUNE = _CONV2D_AUTOTUNE_VOLTA + [
+    triton.Config({**c.kwargs, 'SPATIAL_FAST': True}, num_stages=c.num_stages, num_warps=c.num_warps)
+    for c in _CONV2D_AUTOTUNE_VOLTA
+]
+
+
+@nbx_autotune(
+    configs=maybe_pin_single(_CONV2D_AUTOTUNE, is_conv2d_pinned),
+    key=['batch_dim', 'in_feat_dim', 'in_height', 'in_width_key',
+         'out_feat_dim', 'out_height', 'out_width_key',
+         'kernel_height', 'kernel_width',
+         'stride_height', 'stride_width',
+         'padding_height', 'padding_width',
+         'dilation_height', 'dilation_width',
+         'groups', 'fp16'],
+    cache_results=True,
+)
+@triton.jit
+def conv2d_forward_kernel(
+    input_pointer, weight_pointer, output_pointer,
+    batch_dim, in_feat_dim, in_height, in_width,
+    out_feat_dim, out_height, out_width,
+    in_width_key, out_width_key,
+    input_batch_stride, input_in_feat_stride, input_height_stride, input_width_stride,
+    weight_out_feat_stride, weight_in_feat_stride, weight_height_stride, weight_width_stride,
+    output_batch_stride, output_out_feat_stride, output_height_stride, output_width_stride,
+    kernel_height: tl.constexpr, kernel_width: tl.constexpr,
+    stride_height: tl.constexpr, stride_width: tl.constexpr,
+    padding_height: tl.constexpr, padding_width: tl.constexpr,
+    groups: tl.constexpr, fp16: tl.constexpr,
+    dilation_height: tl.constexpr = 1, dilation_width: tl.constexpr = 1,
+    BLOCK_SIZE_BHW: tl.constexpr = 64,
+    BLOCK_SIZE_INF: tl.constexpr = 32,
+    BLOCK_SIZE_OUTF: tl.constexpr = 64,
+    WIDE_OFFSETS: tl.constexpr = True,
+    SPATIAL_FAST: tl.constexpr = False,
+):
+    # 64-bit offsets only where an operand's extent needs them (`WIDE_OFFSETS`, chosen by the wrapper from the
+    # tensors' extents): 32-bit offsets are exact below 2^31 and the 64-bit ones cost 17 % on Apple's ALUs once the
+    # Metal backend stopped narrowing them (Mac, 2026-10-20, 60 bf16 keys re-proven)
+    if WIDE_OFFSETS:
+        bhw_pid = tl.program_id(0).to(tl.int64)
+        outf_pid = tl.program_id(1).to(tl.int64)
+        group_pid = tl.program_id(2).to(tl.int64)
+    else:
+        bhw_pid = tl.program_id(0)
+        outf_pid = tl.program_id(1)
+        group_pid = tl.program_id(2)
+
+    in_group_dim = in_feat_dim // groups
+    out_group_dim = out_feat_dim // groups
+
+    if SPATIAL_FAST:
+        # The tile is [OUTF, BHW] with the SPATIAL axis fast: NCHW holds a channel plane contiguous along width, so the
+        # input block is loaded [INF, BHW] and the output stored [OUTF, BHW] along it, and the weight block [OUTF, INF]
+        # steps the small kh*kw stride. Loaded [BHW, INF] and stored [BHW, OUTF], both fast axes stepped a whole channel
+        # plane: the Metal backend gathered them element by element — CogVideoX-2b's VAE 3x3 (9,128,242x362) 109 ms
+        # against 55 ms transposed, same values (max |diff| 0), native 31 ms (Mac, 2026-10-21).
+        bhw_offset = bhw_pid * BLOCK_SIZE_BHW + tl.arange(0, BLOCK_SIZE_BHW)
+        bh_offset = bhw_offset // out_width
+        batch_offset = bh_offset // out_height
+
+        outf_offset = outf_pid * BLOCK_SIZE_OUTF + tl.arange(0, BLOCK_SIZE_OUTF)
+        oh_offset = bh_offset % out_height
+        ow_offset = bhw_offset % out_width
+
+        input_pointer += (input_batch_stride * batch_offset +
+                          input_in_feat_stride * group_pid * in_group_dim)[None, :]
+        weight_pointer += (weight_out_feat_stride * outf_offset +
+                           weight_out_feat_stride * group_pid * out_group_dim)[:, None]
+
+        accum = tl.zeros((BLOCK_SIZE_OUTF, BLOCK_SIZE_BHW), dtype=tl.float32)
+
+        for h in range(kernel_height):
+            for w in range(kernel_width):
+                ih_offset = h * dilation_height - padding_height + stride_height * oh_offset
+                iw_offset = w * dilation_width - padding_width + stride_width * ow_offset
+                spatial = input_height_stride * ih_offset + input_width_stride * iw_offset
+                spatial_ok = ((batch_offset < batch_dim) &
+                              (ih_offset >= 0) & (ih_offset < in_height) &
+                              (iw_offset >= 0) & (iw_offset < in_width))
+                for c in range(0, in_group_dim, BLOCK_SIZE_INF):
+                    # 64-bit like every offset above it when the extents need it: `input_in_feat_stride *
+                    # inf_offset` is a product of the channel plane (36 000 000 elements at 6000x6000) and the
+                    # channel index, and past channel 59 it exceeds 2^31 — the 8x upscaler's last conv (64 -> 3, a
+                    # 6344x6344 tile) faulted with an illegal address, four of four, 2026-09-20 (register 58's
+                    # class; the output being small, the band-streaming that saves the 64 -> 64 convs never engaged).
+                    inf_offset = c + tl.arange(0, BLOCK_SIZE_INF)
+                    if WIDE_OFFSETS:
+                        inf_offset = inf_offset.to(tl.int64)
+
+                    curr_inp = (input_pointer +
+                                (input_in_feat_stride * inf_offset)[:, None] +
+                                spatial[None, :])
+                    curr_wt = (weight_pointer +
+                               (weight_in_feat_stride * inf_offset)[None, :] +
+                               (weight_height_stride * h) +
+                               (weight_width_stride * w))
+
+                    inp_mask = (inf_offset < in_group_dim)[:, None] & spatial_ok[None, :]
+                    wt_mask = ((outf_offset < out_group_dim)[:, None] &
+                               (inf_offset < in_group_dim)[None, :])
+
+                    inp_block = tl.load(curr_inp, mask=inp_mask, other=0.0)
+                    wt_block = tl.load(curr_wt, mask=wt_mask, other=0.0)
+
+                    if fp16:
+                        inp_block = inp_block.to(tl.float16)
+                        wt_block = wt_block.to(tl.float16)
+
+                    accum += tl.dot(wt_block, inp_block)
+
+        out_ptr = (output_pointer +
+                   (output_batch_stride * batch_offset)[None, :] +
+                   (output_out_feat_stride * (group_pid * out_group_dim + outf_offset))[:, None] +
+                   (output_height_stride * oh_offset)[None, :] +
+                   (output_width_stride * ow_offset)[None, :])
+        out_mask = ((batch_offset < batch_dim)[None, :] &
+                    (outf_offset < out_group_dim)[:, None] &
+                    (oh_offset < out_height)[None, :] &
+                    (ow_offset < out_width)[None, :])
+
+        tl.store(out_ptr, accum, mask=out_mask)
+    else:
+        bhw_offset = bhw_pid * BLOCK_SIZE_BHW + tl.arange(0, BLOCK_SIZE_BHW)
+        bh_offset = bhw_offset // out_width
+        batch_offset = bh_offset // out_height
+
+        outf_offset = outf_pid * BLOCK_SIZE_OUTF + tl.arange(0, BLOCK_SIZE_OUTF)
+        oh_offset = bh_offset % out_height
+        ow_offset = bhw_offset % out_width
+
+        input_pointer += (input_batch_stride * batch_offset +
+                          input_in_feat_stride * group_pid * in_group_dim)[:, None]
+        weight_pointer += (weight_out_feat_stride * outf_offset +
+                           weight_out_feat_stride * group_pid * out_group_dim)[None, :]
+
+        accum = tl.zeros((BLOCK_SIZE_BHW, BLOCK_SIZE_OUTF), dtype=tl.float32)
+
+        for h in range(kernel_height):
+            for w in range(kernel_width):
+                for c in range(0, in_group_dim, BLOCK_SIZE_INF):
+                    # 64-bit like every offset above it: `input_in_feat_stride *
+                    # inf_offset` is a product of the channel plane (36 000 000
+                    # elements at 6000x6000) and the channel index, and past channel
+                    # 59 it exceeds 2^31 — the 8x upscaler's last conv (64 -> 3, a
+                    # 6344x6344 tile) faulted with an illegal address, four of four,
+                    # 2026-09-20 (register 58's class; the output being small, the
+                    # band-streaming that saves the 64 -> 64 convs never engaged).
+                    inf_offset = c + tl.arange(0, BLOCK_SIZE_INF)
+                    if WIDE_OFFSETS:
+                        inf_offset = inf_offset.to(tl.int64)
+                    ih_offset = h * dilation_height - padding_height + stride_height * oh_offset
+                    iw_offset = w * dilation_width - padding_width + stride_width * ow_offset
+
+                    curr_inp = (input_pointer +
+                                (input_in_feat_stride * inf_offset)[None, :] +
+                                (input_height_stride * ih_offset)[:, None] +
+                                (input_width_stride * iw_offset)[:, None])
+                    curr_wt = (weight_pointer +
+                               (weight_in_feat_stride * inf_offset)[:, None] +
+                               (weight_height_stride * h) +
+                               (weight_width_stride * w))
+
+                    inp_mask = ((batch_offset < batch_dim)[:, None] &
+                                (inf_offset < in_group_dim)[None, :] &
+                                (ih_offset >= 0)[:, None] &
+                                (ih_offset < in_height)[:, None] &
+                                (iw_offset >= 0)[:, None] &
+                                (iw_offset < in_width)[:, None])
+                    wt_mask = ((inf_offset < in_group_dim)[:, None] &
+                               (outf_offset < out_group_dim)[None, :])
+
+                    inp_block = tl.load(curr_inp, mask=inp_mask, other=0.0)
+                    wt_block = tl.load(curr_wt, mask=wt_mask, other=0.0)
+
+                    if fp16:
+                        inp_block = inp_block.to(tl.float16)
+                        wt_block = wt_block.to(tl.float16)
+
+                    accum += tl.dot(inp_block, wt_block)
+
+        out_ptr = (output_pointer +
+                   (output_batch_stride * batch_offset)[:, None] +
+                   (output_out_feat_stride * (group_pid * out_group_dim + outf_offset))[None, :] +
+                   (output_height_stride * oh_offset)[:, None] +
+                   (output_width_stride * ow_offset)[:, None])
+        out_mask = ((batch_offset < batch_dim)[:, None] &
+                    (outf_offset < out_group_dim)[None, :] &
+                    (oh_offset < out_height)[:, None] &
+                    (ow_offset < out_width)[:, None])
+
+        tl.store(out_ptr, accum, mask=out_mask)
+
+
+@triton.jit
+def conv2d_bias_inplace_kernel(out_ptr, bias_ptr, n, HW, C, BLOCK: tl.constexpr):
+    """out[i] += bias[(i // HW) % C], in place, flat over a CONTIGUOUS NCHW
+    output. The conv wrapper's bias epilogue: the out-of-place `add` here
+    allocated a second full-size output per biased convolution — 2450 MB of
+    the measured 8008 MB tile peak on real-esrgan-x8 @1024 (malloc trace,
+    2026-09-21). Elementwise same-offset read/write, so aliasing the input
+    as the output is exact by construction. Fixed launch config — this
+    kernel takes no autotune key and adds nothing to any census."""
+    i = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
+    m = i < n
+    c = (i // HW) % C
+    v = tl.load(out_ptr + i, mask=m)
+    b = tl.load(bias_ptr + c, mask=m)
+    tl.store(out_ptr + i, v + b, mask=m)

nbxtensor/ops/conv2d_m8n8k4.py

@@ -0,0 +1,165 @@
+"""Convolution on the m8n8k4 matrix unit — `conv2d` where the hardware profile declares `matrix_unit`.
+
+An implicit GEMM over the operands in place: per group g, C[m, n] = sum_k A[m, k] B[n, k] with
+m = (batch, oh, ow), n = the group's output channel and k = (c, r, s) flattened in the weight's own order, so a
+contiguous weight [OC, IC/g, KH, KW] is a k-contiguous B read in place (no copy, no im2col buffer) and A is gathered
+from the input through its four strides (padding, stride and dilation are masks and index arithmetic). Same contract
+as `conv2d_forward_kernel` (ops/conv2d.py): fp32
+accumulation (one K tile in the unit, which rounds toward zero, the tiles summed on the FMA units with round to
+nearest — ops/matmul_m8n8k4.py), int64 offsets, every tail masked, the output cast to C's dtype; the bias, when given, is added to the
+fp32 accumulator before that single rounding (the product of two fp16 values is exact in fp32, so the sums are the
+tl.dot kernel's in another order). The tile is the profile's (`matrix_unit.mm`), never tuned at run time: no
+autotune key, nothing to certify.
+
+An operand the DtypeEngine carries SPLIT (SPLIT_A / SPLIT_B: `nbxtensor/dtype/operands.matrix_unit_representation`, the
+profile's `matrix_unit.fp32_split` — the video VAEs' fp32 convolutions) is staged as the GEMM stages it
+(ops/matmul_m8n8k4.py `_stage_split`): hi + lo of the unit's operand dtype under a power-of-two scale per row of A
+(an output position) or column of B (an output channel) of each K tile, three MMAs into two accumulators, the tile's
+sum rescaled and added to the fp32 accumulator on the FMA units.
+
+A is loaded with consecutive lanes on consecutive output rows (consecutive `ow`: coalesced at stride 1), each lane
+holding 4 consecutive k — one 8-byte store into the [k-step x rows, 4] staging of ops/mma_m8n8k4.py.
+"""
+
+import functools
+
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+from .mma_m8n8k4 import acc_rows_cols, mma_k4, stage_k_inner
+from .matmul_m8n8k4 import _bcast_cols, _bcast_rows, _stage_split, mm_layouts
+
+
+@gluon.jit
+def conv2d_m8n8k4_kernel(
+    X, Wt, Bias, C,
+    M, NG, K, IH, IW, OH, OW,
+    stride_xn, stride_xc, stride_xh, stride_xw,
+    stride_wo, stride_cn, stride_cc, stride_ch, stride_cw,
+    CG, OG,
+    KH: gl.constexpr, KW: gl.constexpr, SH_: gl.constexpr, SW_: gl.constexpr, PH: gl.constexpr,
+    PW: gl.constexpr, DH: gl.constexpr, DW: gl.constexpr,
+    HAS_BIAS: gl.constexpr,
+    SPLIT_A: gl.constexpr, SPLIT_B: gl.constexpr, HI_EXP: gl.constexpr, LO_SHIFT: gl.constexpr,
+    MB: gl.constexpr, NB: gl.constexpr, BK: gl.constexpr, GROUP_M: gl.constexpr,
+    L: gl.constexpr, LA: gl.constexpr, LB: gl.constexpr, GA: gl.constexpr, GB: gl.constexpr,
+):
+    BM: gl.constexpr = 16 * MB
+    BN: gl.constexpr = 16 * NB
+    J: gl.constexpr = BK // 4
+    KHW: gl.constexpr = KH * KW
+    SH: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, order=[1, 0])
+
+    pid = gl.program_id(0).to(gl.int64)
+    g = gl.program_id(1).to(gl.int64)                      # the group
+    X = X + g * CG * stride_xc
+    Wt = Wt + g * OG * stride_wo
+    C = C + g * OG * stride_cc
+    num_pid_m = gl.cdiv(M, BM)
+    num_pid_n = gl.cdiv(NG, BN)
+    num_pid_in_group = GROUP_M * num_pid_n
+    first_pid_m = (pid // num_pid_in_group) * GROUP_M
+    group_size_m = gl.minimum(num_pid_m - first_pid_m, GROUP_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    m0 = pid_m * BM
+    n0 = pid_n * BN
+
+    sa = gl.allocate_shared_memory(gl.float16, [J * BM, 4], SH)
+    sb = gl.allocate_shared_memory(gl.float16, [J * BN, 4], SH)
+    if SPLIT_A:
+        sal = gl.allocate_shared_memory(gl.float16, [J * BM, 4], SH)
+    if SPLIT_B:
+        sbl = gl.allocate_shared_memory(gl.float16, [J * BN, 4], SH)
+
+    # A rows: (batch, oh, ow) once per tile; the k columns are decoded per K tile against constexpr divisors.
+    ma1 = m0 + gl.arange(0, BM, layout=gl.SliceLayout(1, GA))
+    ow_ = ma1 % OW
+    oh_ = (ma1 // OW) % OH
+    xrow = ((ma1 // OW) // OH) * stride_xn
+    ih0 = (oh_ * SH_ - PH)[:, None]
+    iw0 = (ow_ * SW_ - PW)[:, None]
+    ma = (ma1 < M)[:, None]
+    xrow = xrow[:, None]
+    ka1 = gl.arange(0, BK, layout=gl.SliceLayout(0, GA))
+    # B [BN, BK], k contiguous: the weight read in place.
+    nb_ = gl.arange(0, BN, layout=gl.SliceLayout(1, GB))[:, None]
+    kb = gl.arange(0, BK, layout=gl.SliceLayout(0, GB))[None, :]
+    pb = Wt + (n0 + nb_) * stride_wo + kb
+    mb_ = (n0 + nb_) < NG
+
+    acc = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+    ta = _gather(X, xrow, ih0, iw0, ma, ka1, 0, K, IH, IW, stride_xc, stride_xh, stride_xw, KHW, KW, DH, DW)
+    tb = gl.load(pb, mask=mb_ & (kb < K), other=0.0)
+    for k0 in range(0, K, BK):
+        gl.barrier()                                       # every warp is done reading the previous tile
+        if SPLIT_A:
+            inv_a = _stage_split(sa, sal, ta, True, BM, BK, HI_EXP, LO_SHIFT)
+        else:
+            stage_k_inner(sa, ta, BM, BK)
+        if SPLIT_B:
+            inv_b = _stage_split(sb, sbl, tb, True, BN, BK, HI_EXP, LO_SHIFT)
+        else:
+            stage_k_inner(sb, tb, BN, BK)
+        gl.barrier()
+        if k0 + BK < K:                                    # the next tile loads while the HMMAs run
+            ta = _gather(X, xrow, ih0, iw0, ma, ka1, k0 + BK, K, IH, IW, stride_xc, stride_xh, stride_xw, KHW, KW,
+                         DH, DW)
+            tb = gl.load(pb + (k0 + BK), mask=mb_ & (kb + k0 + BK < K), other=0.0)
+        if SPLIT_A or SPLIT_B:
+            t_hh = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+            t_lo = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+            for j in gl.static_range(J):
+                t_hh = mma_k4(sa, j * BM, sb, j * BN, t_hh, MB, NB, LA, LB, L)
+                if SPLIT_A:
+                    t_lo = mma_k4(sal, j * BM, sb, j * BN, t_lo, MB, NB, LA, LB, L)
+                if SPLIT_B:
+                    t_lo = mma_k4(sa, j * BM, sbl, j * BN, t_lo, MB, NB, LA, LB, L)
+            t = t_hh + t_lo * (1.0 / (1 << LO_SHIFT))
+            if SPLIT_A:
+                t = t * _bcast_rows(gl.reshape(inv_a, [MB, 2, 8]), L)
+            if SPLIT_B:
+                t = t * _bcast_cols(gl.reshape(inv_b, [NB, 2, 2, 4]), L)
+            acc = acc + t
+        else:
+            t = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+            for j in gl.static_range(J):
+                t = mma_k4(sa, j * BM, sb, j * BN, t, MB, NB, LA, LB, L)
+            acc = acc + t                                  # across tiles on the FMA units, round to nearest
+
+    rows, cols = acc_rows_cols(L, MB, NB)
+    rows = m0 + rows
+    cols = n0 + cols
+    mask = (rows < M) & (cols < NG)
+    if HAS_BIAS:
+        bias = gl.load(Bias + g * OG + cols.to(gl.int64), mask=cols < NG, other=0.0)
+        acc = acc + bias.to(gl.float32)
+    r64 = rows.to(gl.int64)
+    ow_o = r64 % OW
+    oh_o = (r64 // OW) % OH
+    n_o = (r64 // OW) // OH
+    gl.store(C + n_o * stride_cn + oh_o * stride_ch + ow_o * stride_cw + cols.to(gl.int64) * stride_cc,
+             acc.to(C.dtype.element_ty), mask=mask)
+
+
+@gluon.jit
+def _gather(X, xrow, ih0, iw0, ma, ka1, k0, K, IH, IW, stride_xc, stride_xh, stride_xw,
+            KHW: gl.constexpr, KW: gl.constexpr, DH: gl.constexpr, DW: gl.constexpr):
+    """A[rows, k0:k0+BK]: input[batch, c, oh*sh - ph + r*dh, ow*sw - pw + s*dw] for k = (c, r, s), 0 outside."""
+    k = k0 + ka1
+    c = k // KHW
+    rs = k % KHW
+    ih = ih0 + ((rs // KW) * DH)[None, :]
+    iw = iw0 + ((rs % KW) * DW)[None, :]
+    m = ma & (k < K)[None, :] & (ih >= 0) & (ih < IH) & (iw >= 0) & (iw < IW)
+    p = X + xrow + (c.to(gl.int64) * stride_xc)[None, :] + ih.to(gl.int64) * stride_xh + iw.to(gl.int64) * stride_xw
+    return gl.load(p, mask=m, other=0.0)
+
+
+@functools.lru_cache(maxsize=None)
+def conv_layouts(MB: int, NB: int, WM: int, WN: int, BK: int) -> dict:
+    """The GEMM's layouts (ops/matmul_m8n8k4.py, B k-contiguous) with the gathered A walked row-fastest: 32
+    consecutive output rows per warp, 4 consecutive k per lane."""
+    lay = dict(mm_layouts(MB, NB, WM, WN, BK, True))
+    lay["GA"] = gl.BlockedLayout([1, 4], [32, 1], [WM * WN, 1], [0, 1])
+    return lay

nbxtensor/ops/conv_depthwise2d.py

@@ -0,0 +1,154 @@
+"""Depthwise 2D convolution — pure @triton.jit kernel.
+
+Depthwise convolution: each input channel is convolved independently with
+its own filter (groups = C_in). Output channels = C_in * channel_multiplier
+(typically channel_multiplier = 1).
+
+Input:  [N, C, IH, IW]
+Weight: [C_out, 1, KH, KW]  (groups = C_in)
+Output: [N, C_out, OH, OW]
+
+OH = (IH + 2*pad_h - dilation_h*(KH-1) - 1) // stride_h + 1
+OW = (IW + 2*pad_w - dilation_w*(KW-1) - 1) // stride_w + 1
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def conv_depthwise2d_kernel(
+    input_ptr,
+    weight_ptr,
+    output_ptr,
+    # Input shape
+    N, C_in, IH, IW,
+    # Weight shape: [C_out, 1, KH, KW]
+    C_out, KH, KW,
+    # Output shape
+    OH, OW,
+    # Conv parameters
+    stride_h, stride_w,
+    pad_h, pad_w,
+    dilation_h, dilation_w,
+    # Channel multiplier: C_out // C_in (usually 1)
+    ch_mult,
+    BLOCK_X: tl.constexpr,  # OW tile
+    BLOCK_Y: tl.constexpr,  # OH tile
+):
+    """Depthwise conv2d kernel.
+
+    Grid: (cdiv(OW, BLOCK_X), cdiv(OH, BLOCK_Y), N * C_out)
+    Each program computes a BLOCK_Y x BLOCK_X tile of one output channel.
+    """
+    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)
+
+    n = pid_nc // C_out
+    c_out = pid_nc % C_out
+    c_in = c_out // ch_mult
+
+    ow = pid_x * BLOCK_X + tl.arange(0, BLOCK_X)  # [BLOCK_X]
+    oh = pid_y * BLOCK_Y + tl.arange(0, BLOCK_Y)  # [BLOCK_Y]
+
+    mask_h = oh < OH
+    mask_w = ow < OW
+    mask_2d = mask_h[:, None] & mask_w[None, :]
+
+    acc = tl.zeros([BLOCK_Y, BLOCK_X], dtype=tl.float32)
+
+    # Base pointer for input channel
+    in_base = n * (C_in * IH * IW) + c_in * (IH * IW)
+
+    # Base pointer for weight
+    w_base = c_out * (KH * KW)  # weight shape is [C_out, 1, KH, KW], 1 is implicit
+
+    for kh in range(0, KH):
+        for kw in range(0, KW):
+            # Input coordinates
+            ih = oh[:, None] * stride_h - pad_h + kh * dilation_h  # [BLOCK_Y, 1]
+            iw = ow[None, :] * stride_w - pad_w + kw * dilation_w  # [1, BLOCK_X]
+
+            # Bounds check
+            valid_h = (ih >= 0) & (ih < IH)
+            valid_w = (iw >= 0) & (iw < IW)
+            valid = mask_2d & valid_h & valid_w
+
+            # Load input value
+            in_offset = in_base + ih * IW + iw
+            x_val = tl.load(input_ptr + in_offset, mask=valid, other=0.0).to(tl.float32)
+
+            # Load weight value (scalar per kh, kw)
+            w_offset = w_base + kh * KW + kw
+            w_val = tl.load(weight_ptr + w_offset).to(tl.float32)
+
+            acc += x_val * w_val
+
+    # Store output
+    out_base = n * (C_out * OH * OW) + c_out * (OH * OW)
+    out_offset = out_base + oh[:, None] * OW + ow[None, :]
+    tl.store(output_ptr + out_offset, acc, mask=mask_2d)
+
+
+@triton.jit
+def conv_depthwise2d_bias_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    N, C_in, IH, IW,
+    C_out, KH, KW,
+    OH, OW,
+    stride_h, stride_w,
+    pad_h, pad_w,
+    dilation_h, dilation_w,
+    ch_mult,
+    BLOCK_X: tl.constexpr,
+    BLOCK_Y: tl.constexpr,
+):
+    """Depthwise conv2d with bias, same structure as above."""
+    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)
+
+    n = pid_nc // C_out
+    c_out = pid_nc % C_out
+    c_in = c_out // ch_mult
+
+    ow = pid_x * BLOCK_X + tl.arange(0, BLOCK_X)
+    oh = pid_y * BLOCK_Y + tl.arange(0, BLOCK_Y)
+
+    mask_h = oh < OH
+    mask_w = ow < OW
+    mask_2d = mask_h[:, None] & mask_w[None, :]
+
+    acc = tl.zeros([BLOCK_Y, BLOCK_X], dtype=tl.float32)
+
+    in_base = n * (C_in * IH * IW) + c_in * (IH * IW)
+    w_base = c_out * (KH * KW)
+
+    for kh in range(0, KH):
+        for kw in range(0, KW):
+            ih = oh[:, None] * stride_h - pad_h + kh * dilation_h
+            iw = ow[None, :] * stride_w - pad_w + kw * dilation_w
+
+            valid_h = (ih >= 0) & (ih < IH)
+            valid_w = (iw >= 0) & (iw < IW)
+            valid = mask_2d & valid_h & valid_w
+
+            in_offset = in_base + ih * IW + iw
+            x_val = tl.load(input_ptr + in_offset, mask=valid, other=0.0).to(tl.float32)
+
+            w_offset = w_base + kh * KW + kw
+            w_val = tl.load(weight_ptr + w_offset).to(tl.float32)
+
+            acc += x_val * w_val
+
+    # Add bias
+    bias_val = tl.load(bias_ptr + c_out).to(tl.float32)
+    acc += bias_val
+
+    out_base = n * (C_out * OH * OW) + c_out * (OH * OW)
+    out_offset = out_base + oh[:, None] * OW + ow[None, :]
+    tl.store(output_ptr + out_offset, acc, mask=mask_2d)