nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/max_pool2d.py

@@ -0,0 +1,74 @@
+"""Max pooling 2D with indices — pure @triton.jit kernel.
+
+Extracted from FlagGems max_pool2d_with_indices.py. Forward only (inference).
+2D blocked approach, returns both output values and indices of max elements.
+Uses float('-inf') as min value instead of FlagGems get_dtype_min().
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def max_pool2d_forward_kernel(
+    input_ptr,
+    output_ptr,
+    indices_ptr,
+    in_stride_n, in_stride_c, in_stride_h, in_stride_w,
+    in_c, in_h, in_w,
+    out_h, out_w,
+    kernel_h: tl.constexpr,
+    kernel_w: tl.constexpr,
+    stride_h: tl.constexpr,
+    stride_w: tl.constexpr,
+    padding_h: tl.constexpr,
+    padding_w: tl.constexpr,
+    dilation_h: tl.constexpr,
+    dilation_w: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_W: tl.constexpr,
+):
+    pid_nc = tl.program_id(0).to(tl.int64)
+    pid_hw = tl.program_id(1).to(tl.int64)
+    num_w_blocks = tl.cdiv(out_w, BLOCK_W)
+    h_block_idx = pid_hw // num_w_blocks
+    w_block_idx = pid_hw % num_w_blocks
+    n_idx = pid_nc // in_c
+    c_idx = pid_nc % in_c
+
+    h_out_offsets = h_block_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    w_out_offsets = w_block_idx * BLOCK_W + tl.arange(0, BLOCK_W)
+
+    min_val = float('-inf')
+    max_val_acc = tl.full((BLOCK_H, BLOCK_W), min_val, dtype=input_ptr.type.element_ty)
+    max_idx_acc = tl.full((BLOCK_H, BLOCK_W), -1, dtype=tl.int64)
+
+    input_base_ptr = input_ptr + n_idx * in_stride_n + c_idx * in_stride_c
+
+    for kh in tl.static_range(0, kernel_h):
+        for kw in tl.static_range(0, kernel_w):
+            h_in = h_out_offsets[:, None] * stride_h - padding_h + kh * dilation_h
+            w_in = w_out_offsets[None, :] * stride_w - padding_w + kw * dilation_w
+            in_mask = (h_in >= 0) & (h_in < in_h) & (w_in >= 0) & (w_in < in_w)
+            input_offset = h_in * in_stride_h + w_in * in_stride_w
+            current_val = tl.load(
+                input_base_ptr + input_offset, mask=in_mask, other=min_val
+            )
+            current_idx = h_in * in_w + w_in
+
+            is_new_max = current_val > max_val_acc
+            max_val_acc = tl.where(is_new_max, current_val, max_val_acc)
+            max_idx_acc = tl.where(is_new_max & in_mask, current_idx, max_idx_acc)
+
+    out_base_ptr = output_ptr + pid_nc * out_h * out_w
+    indices_base_ptr = indices_ptr + pid_nc * out_h * out_w
+    output_block_ptr = (
+        out_base_ptr + h_out_offsets[:, None] * out_w + w_out_offsets[None, :]
+    )
+    indices_block_ptr = (
+        indices_base_ptr + h_out_offsets[:, None] * out_w + w_out_offsets[None, :]
+    )
+
+    out_mask = (h_out_offsets[:, None] < out_h) & (w_out_offsets[None, :] < out_w)
+    tl.store(output_block_ptr, max_val_acc, mask=out_mask)
+    tl.store(indices_block_ptr, max_idx_acc, mask=out_mask)

nbxtensor/ops/max_pool3d.py

@@ -0,0 +1,93 @@
+"""Max pooling 3D with indices — pure @triton.jit kernel.
+
+3D (temporal) extension of ops/max_pool2d.py (FlagGems-derived blocked
+kernel): one program per (n, c, t_out) output plane block, with a
+static kt x kh x kw window scan. Forward only (inference). Returns both
+output values and int64 indices of the max elements; indices follow the
+PyTorch max_pool3d_with_indices convention — flattened offset into the
+(T, H, W) volume of each (n, c) plane, ties resolved to the FIRST
+maximum in (t, h, w) window scan order (strict `>` update, same as the
+2D kernel). Uses float('-inf') as the identity like the 2D kernel.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def max_pool3d_forward_kernel(
+    input_ptr,
+    output_ptr,
+    indices_ptr,
+    in_stride_n, in_stride_c, in_stride_t, in_stride_h, in_stride_w,
+    in_c, in_t, in_h, in_w,
+    out_t, out_h, out_w,
+    kernel_t: tl.constexpr,
+    kernel_h: tl.constexpr,
+    kernel_w: tl.constexpr,
+    stride_t: tl.constexpr,
+    stride_h: tl.constexpr,
+    stride_w: tl.constexpr,
+    padding_t: tl.constexpr,
+    padding_h: tl.constexpr,
+    padding_w: tl.constexpr,
+    dilation_t: tl.constexpr,
+    dilation_h: tl.constexpr,
+    dilation_w: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    BLOCK_W: tl.constexpr,
+):
+    pid_nct = tl.program_id(0).to(tl.int64)
+    pid_hw = tl.program_id(1).to(tl.int64)
+    num_w_blocks = tl.cdiv(out_w, BLOCK_W)
+    h_block_idx = pid_hw // num_w_blocks
+    w_block_idx = pid_hw % num_w_blocks
+    t_out_idx = pid_nct % out_t
+    nc_idx = pid_nct // out_t
+    n_idx = nc_idx // in_c
+    c_idx = nc_idx % in_c
+
+    h_out_offsets = h_block_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    w_out_offsets = w_block_idx * BLOCK_W + tl.arange(0, BLOCK_W)
+
+    min_val = float('-inf')
+    max_val_acc = tl.full((BLOCK_H, BLOCK_W), min_val, dtype=input_ptr.type.element_ty)
+    max_idx_acc = tl.full((BLOCK_H, BLOCK_W), -1, dtype=tl.int64)
+
+    input_base_ptr = input_ptr + n_idx * in_stride_n + c_idx * in_stride_c
+
+    for kt in tl.static_range(0, kernel_t):
+        t_in = t_out_idx * stride_t - padding_t + kt * dilation_t
+        t_ok = (t_in >= 0) & (t_in < in_t)
+        for kh in tl.static_range(0, kernel_h):
+            for kw in tl.static_range(0, kernel_w):
+                h_in = h_out_offsets[:, None] * stride_h - padding_h + kh * dilation_h
+                w_in = w_out_offsets[None, :] * stride_w - padding_w + kw * dilation_w
+                in_mask = (t_ok
+                           & (h_in >= 0) & (h_in < in_h)
+                           & (w_in >= 0) & (w_in < in_w))
+                input_offset = (t_in * in_stride_t
+                                + h_in * in_stride_h + w_in * in_stride_w)
+                current_val = tl.load(
+                    input_base_ptr + input_offset, mask=in_mask, other=min_val
+                )
+                # PyTorch convention: flattened (t, h, w) offset per (n, c).
+                current_idx = ((t_in * in_h + h_in) * in_w + w_in).to(tl.int64)
+
+                is_new_max = current_val > max_val_acc
+                max_val_acc = tl.where(is_new_max, current_val, max_val_acc)
+                max_idx_acc = tl.where(is_new_max & in_mask, current_idx, max_idx_acc)
+
+    plane = out_h * out_w
+    out_base_ptr = output_ptr + pid_nct * plane
+    indices_base_ptr = indices_ptr + pid_nct * plane
+    output_block_ptr = (
+        out_base_ptr + h_out_offsets[:, None] * out_w + w_out_offsets[None, :]
+    )
+    indices_block_ptr = (
+        indices_base_ptr + h_out_offsets[:, None] * out_w + w_out_offsets[None, :]
+    )
+
+    out_mask = (h_out_offsets[:, None] < out_h) & (w_out_offsets[None, :] < out_w)
+    tl.store(output_block_ptr, max_val_acc, mask=out_mask)
+    tl.store(indices_block_ptr, max_idx_acc, mask=out_mask)

nbxtensor/ops/max_reduce.py

@@ -0,0 +1,91 @@
+"""Max reduction — pure @triton.jit kernels.
+
+Global two-pass max and dim-specific max (returns value + index).
+Extracted from FlagGems (Apache-2.0).
+
+Note: FlagGems uses `get_dtype_min` from utils.limits. We inline the
+sentinel via float('-inf') which Triton handles for all float dtypes.
+For the dim kernel, bf16 accumulation is promoted to fp32.
+"""
+
+import triton
+import triton.language as tl
+
+
+# ---------------------------------------------------------------------------
+# Global reduction: two-pass max
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def max_kernel_mid(
+    inp,
+    mid,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pass 1: compute block-local max."""
+    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)
+    inp_ptrs = inp + offset
+    mask = offset < M
+    inp_val = tl.load(inp_ptrs, mask=mask, other=float('-inf'))
+    max_val = tl.max(inp_val)
+    mid_ptr = mid + pid
+    tl.store(mid_ptr, max_val)
+
+
+@triton.jit
+def max_kernel_result(mid, out, mid_size, BLOCK_MID: tl.constexpr):
+    """Pass 2: reduce partial maxes into scalar."""
+    offset = tl.arange(0, BLOCK_MID)
+    mid_ptrs = mid + offset
+    mask = offset < mid_size
+    mid_val = tl.load(mid_ptrs, mask=mask, other=float('-inf'))
+    max_val = tl.max(mid_val)
+    tl.store(out, max_val)
+
+
+# ---------------------------------------------------------------------------
+# Dim-specific: max along one axis (returns value + argmax index)
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def max_kernel(
+    inp,
+    out_value,
+    out_index,
+    M,
+    N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Max reduction along compressed inner dim N.
+
+    inp: [M, N] (dim-compressed) → out_value: [M], out_index: [M] (int64)
+    Tracks both maximum value and its index (argmax).
+    bf16 inputs are accumulated in fp32 to avoid precision issues.
+    """
+    pid_m = tl.program_id(0).to(tl.int64)
+    m_offset = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    dtype = inp.type.element_ty
+    acc_type = tl.float32 if dtype is tl.bfloat16 else dtype
+    result_value = tl.full([BLOCK_M], value=float('-inf'), dtype=acc_type)
+    result_index = tl.zeros([BLOCK_M], dtype=tl.int64)
+    for i in range(0, N, BLOCK_N):
+        n_offset = i + tl.arange(0, BLOCK_N)
+        offset = m_offset[:, None] * N + n_offset[None, :]
+        mask = (m_offset[:, None] < M) & (n_offset[None, :] < N)
+        inp_ptrs = inp + offset
+        inp_vals = tl.load(inp_ptrs, mask=mask, other=float('-inf'))
+        max_value, max_index = tl.max(inp_vals, axis=1, return_indices=True)
+        update_mask = max_value > result_value
+        result_value = tl.where(update_mask, max_value, result_value)
+        result_index = tl.where(update_mask, i + max_index, result_index)
+
+    mask1 = m_offset < M
+    offset_index = m_offset
+    out_value_ptrs = out_value + offset_index
+    out_index_ptrs = out_index + offset_index
+    tl.store(out_value_ptrs, result_value, mask=mask1)
+    tl.store(out_index_ptrs, result_index, mask=mask1)

nbxtensor/ops/maximum.py

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

nbxtensor/ops/mean.py

@@ -0,0 +1,40 @@
+"""Mean 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 mean_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,
+):
+    """Mean 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 / feat_dim, mask=batch_mask)

nbxtensor/ops/min_reduce.py

@@ -0,0 +1,91 @@
+"""Min reduction — pure @triton.jit kernels.
+
+Global two-pass min and dim-specific min (returns value + index).
+Extracted from FlagGems (Apache-2.0).
+
+Note: FlagGems uses `get_dtype_max` from utils.limits. We inline the
+sentinel via float('inf') which Triton handles for all float dtypes.
+For the dim kernel, bf16 accumulation is promoted to fp32.
+"""
+
+import triton
+import triton.language as tl
+
+
+# ---------------------------------------------------------------------------
+# Global reduction: two-pass min
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def min_kernel_mid(
+    inp,
+    mid,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pass 1: compute block-local min."""
+    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)
+    inp_ptrs = inp + offset
+    mask = offset < M
+    inp_val = tl.load(inp_ptrs, mask=mask, other=float('inf'))
+    min_val = tl.min(inp_val)
+    mid_ptr = mid + pid
+    tl.store(mid_ptr, min_val)
+
+
+@triton.jit
+def min_kernel_result(mid, out, mid_size, BLOCK_MID: tl.constexpr):
+    """Pass 2: reduce partial mins into scalar."""
+    offset = tl.arange(0, BLOCK_MID)
+    mid_ptrs = mid + offset
+    mask = offset < mid_size
+    mid_val = tl.load(mid_ptrs, mask=mask, other=float('inf'))
+    min_val = tl.min(mid_val)
+    tl.store(out, min_val)
+
+
+# ---------------------------------------------------------------------------
+# Dim-specific: min along one axis (returns value + argmin index)
+# ---------------------------------------------------------------------------
+
+@triton.jit
+def min_kernel(
+    inp,
+    out_value,
+    out_index,
+    M,
+    N,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    """Min reduction along compressed inner dim N.
+
+    inp: [M, N] (dim-compressed) → out_value: [M], out_index: [M] (int64)
+    Tracks both minimum value and its index (argmin).
+    bf16 inputs are accumulated in fp32 to avoid precision issues.
+    """
+    pid_m = tl.program_id(0).to(tl.int64)
+    m_offset = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+
+    dtype = inp.type.element_ty
+    acc_type = tl.float32 if dtype is tl.bfloat16 else dtype
+    min_values = tl.full([BLOCK_M], dtype=acc_type, value=float('inf'))
+    argmin_values = tl.full([BLOCK_M], dtype=tl.int64, value=0)
+    for start_n in range(0, N, BLOCK_N):
+        n_offset = start_n + tl.arange(0, BLOCK_N)
+        offset = m_offset[:, None] * N + n_offset[None, :]
+        mask = (m_offset[:, None] < M) & (n_offset[None, :] < N)
+        inp_ptrs = inp + offset
+        inp_vals = tl.load(inp_ptrs, mask=mask, other=float('inf'))
+        local_min, local_argmin = tl.min(inp_vals, 1, return_indices=True)
+        update = local_min < min_values
+        min_values = tl.where(update, local_min, min_values)
+        argmin_values = tl.where(update, start_n + local_argmin, argmin_values)
+
+    offset_index = m_offset
+    out_value_ptrs = out_value + offset_index
+    out_index_ptrs = out_index + offset_index
+    mask1 = m_offset < M
+    tl.store(out_value_ptrs, min_values, mask=mask1)
+    tl.store(out_index_ptrs, argmin_values, mask=mask1)

nbxtensor/ops/minimum.py

@@ -0,0 +1,18 @@
+"""Element-wise minimum — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def minimum_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)
+    tl.store(output_ptr + offset, tl.minimum(x, y), mask=mask)

nbxtensor/ops/mish.py

@@ -0,0 +1,22 @@
+"""Mish — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+from ._common import tanh_fn
+
+@triton.jit
+def mish_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)
+    # mish = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x)))
+    out = x_fp32 * tanh_fn(tl.log(1.0 + tl.exp(x_fp32)))
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/mma_m8n8k4.py

@@ -0,0 +1,152 @@
+"""The m8n8k4 matrix unit (sm_70 tensor cores, SASS HMMA.884) as a Gluon building block.
+
+Triton >= 3.3 lowers `tl.dot` to scalar FMA below sm_80 (triton-lang/triton#5066), so an arch whose profile
+declares `matrix_unit.shape: m8n8k4` gets its tensor cores from here: our own Gluon kernels issue
+`mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32` through `inline_asm_elementwise`, with the operands placed in
+registers by linear layouts. Method and proofs: docs/reference/tensor-cores-through-linear-layouts.md.
+
+The fragment (CUTLASS cute/atom/mma_traits_sm70.hpp, SM70_8x8x4_F32F16F16F32_TN): lane = a0 + 2 a1 + 4 q0 + 8 q1
++ 16 b, thread t = a0 + 2 a1 + 4 b of quadpair q holds A row t and B column t (k = 0..3), and accumulator
+register v = v0 + 2 v1 + 4 v2 holds D[m = a0 + 2 v1 + 4 b, n = v0 + 2 a1 + 4 v2]; the four quadpairs tile 16x16.
+
+A [BM, BN] accumulator lives as the 7-D tensor [MB, 2, 8, NB, 2, 2, 4] = (mb, mi3, mi2..0, nb, ni3, ni2, ni1..0),
+m = 16 mb + mi, n = 16 nb + ni (row-major order, so a reshape to [BM, BN] moves nothing). In these coordinates the
+A fragment depends only on (mb, mi3, mi2..0, ni1..0) and the B fragment only on (mi2..0, nb, ni3, ni1..0): each is
+a plain broadcast of [MB, 2, 8, 1, 1, 1, 4] / [1, 1, 8, NB, 2, 1, 4], which `inline_asm_elementwise` broadcasts
+itself, so one register fragment feeds every sub-tile that shares it. Operands come from shared memory laid out
+[k-step x rows, 4] (k innermost, 8 bytes per row) and are read in a linear layout derived from the accumulator's;
+the reshape/permute to the broadcast shape is a relabelling, which `convert_layout(assert_trivial=True)` proves.
+"""
+
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+#: What the instruction below takes, as a hardware profile's `matrix_unit` names it (checked by
+#: `_configs.matrix_unit`: a profile cannot declare dtypes this instruction does not compute in). `k` is the
+#: columns one instruction reduces: the split flash's `block_k` is a multiple of it (KSA/KSB count instructions).
+UNIT = {"shape": "m8n8k4", "operand_dtype": "float16", "accumulator_dtype": "float32", "k": 4}
+
+MMA_ASM: gl.constexpr = gl.constexpr(
+    "mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 {$0,$1,$2,$3,$4,$5,$6,$7}, {$8, $9}, {$12, $13},"
+    " {$16,$17,$18,$19,$20,$21,$22,$23};")
+# pack=8: eight accumulator elements per call; f16 operands arrive two per 32-bit register, so A and B are four
+# registers each of which the instruction reads the first two (k = 0..3; the other two repeat them).
+MMA_CONSTRAINTS: gl.constexpr = gl.constexpr(",".join(["=f"] * 8 + ["r"] * 8 + ["f"] * 8))
+
+
+def _bits(n: int):
+    return [1 << i for i in range(n.bit_length() - 1)]
+
+
+def _e7(d: int, x: int):
+    v = [0] * 7
+    v[d] = x
+    return v
+
+
+def acc_layout(MB: int, NB: int, WM: int, WN: int):
+    """The [MB, 2, 8, NB, 2, 2, 4] accumulator over WM x WN warps (WM | MB, WN | NB, all powers of two)."""
+    rm = [x * WM for x in _bits(MB // WM)]
+    rn = [x * WN for x in _bits(NB // WN)]
+    return gl.DistributedLinearLayout(
+        reg_bases=[_e7(6, 1), _e7(2, 2), _e7(5, 1)] + [_e7(3, x) for x in rn] + [_e7(0, x) for x in rm],
+        lane_bases=[_e7(2, 1), _e7(6, 2), _e7(1, 1), _e7(4, 1), _e7(2, 4)],
+        warp_bases=[_e7(0, x) for x in _bits(WM)] + [_e7(3, x) for x in _bits(WN)],
+        block_bases=[], shape=[MB, 2, 8, NB, 2, 2, 4])
+
+
+def frag_a_layout(MB: int, NB: int, WM: int, WN: int):
+    """One k-step of the A operand, [16 MB, 4] = (row, k), in the registers acc_layout(MB, NB, WM, WN) needs."""
+    rm = [x * WM for x in _bits(MB // WM)]
+    rn = [x * WN for x in _bits(NB // WN)]
+    return gl.DistributedLinearLayout(
+        reg_bases=[[0, 1], [0, 2], [0, 0]] + [[0, 0]] * len(rn) + [[16 * x, 0] for x in rm],
+        lane_bases=[[1, 0], [2, 0], [8, 0], [0, 0], [4, 0]],
+        warp_bases=[[16 * x, 0] for x in _bits(WM)] + [[0, 0]] * len(_bits(WN)),
+        block_bases=[], shape=[16 * MB, 4])
+
+
+def frag_b_layout(MB: int, NB: int, WM: int, WN: int):
+    """One k-step of the B operand, [16 NB, 4] = (column, k), in the registers acc_layout(MB, NB, WM, WN) needs."""
+    rm = [x * WM for x in _bits(MB // WM)]
+    rn = [x * WN for x in _bits(NB // WN)]
+    return gl.DistributedLinearLayout(
+        reg_bases=[[0, 1], [0, 2], [0, 0]] + [[16 * x, 0] for x in rn] + [[0, 0]] * len(rm),
+        lane_bases=[[1, 0], [2, 0], [0, 0], [8, 0], [4, 0]],
+        warp_bases=[[0, 0]] * len(_bits(WM)) + [[16 * x, 0] for x in _bits(WN)],
+        block_bases=[], shape=[16 * NB, 4])
+
+
+@gluon.jit
+def acc_rows_cols(L: gl.constexpr, MB: gl.constexpr, NB: gl.constexpr):
+    """Row and column index of every accumulator element, broadcastable over L ([.., 1, 1, 1, 1] / [1, 1, 1, ..])."""
+    o0 = gl.arange(0, MB, layout=gl.SliceLayout(1, gl.SliceLayout(2, gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(5, gl.SliceLayout(6, L)))))))
+    o1 = gl.arange(0, 2, layout=gl.SliceLayout(0, gl.SliceLayout(2, gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(5, gl.SliceLayout(6, L)))))))
+    o2 = gl.arange(0, 8, layout=gl.SliceLayout(0, gl.SliceLayout(1, gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(5, gl.SliceLayout(6, L)))))))
+    o3 = gl.arange(0, NB, layout=gl.SliceLayout(0, gl.SliceLayout(1, gl.SliceLayout(2, gl.SliceLayout(4, gl.SliceLayout(5, gl.SliceLayout(6, L)))))))
+    o4 = gl.arange(0, 2, layout=gl.SliceLayout(0, gl.SliceLayout(1, gl.SliceLayout(2, gl.SliceLayout(3, gl.SliceLayout(5, gl.SliceLayout(6, L)))))))
+    o5 = gl.arange(0, 2, layout=gl.SliceLayout(0, gl.SliceLayout(1, gl.SliceLayout(2, gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(6, L)))))))
+    o6 = gl.arange(0, 4, layout=gl.SliceLayout(0, gl.SliceLayout(1, gl.SliceLayout(2, gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(5, L)))))))
+    rows = o0[:, None, None, None, None, None, None] * 16 + o1[None, :, None, None, None, None, None] * 8 \
+        + o2[None, None, :, None, None, None, None]
+    cols = o3[None, None, None, :, None, None, None] * 16 + o4[None, None, None, None, :, None, None] * 8 \
+        + o5[None, None, None, None, None, :, None] * 4 + o6[None, None, None, None, None, None, :]
+    return rows, cols
+
+
+@gluon.jit
+def row_layout_rows(L: gl.constexpr, MB: gl.constexpr):
+    """Row index [MB, 2, 8] in the layout of a row statistic of L (L with its four column dims reduced)."""
+    RS: gl.constexpr = gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(5, gl.SliceLayout(6, L))))
+    r0 = gl.arange(0, MB, layout=gl.SliceLayout(1, gl.SliceLayout(2, RS)))
+    r1 = gl.arange(0, 2, layout=gl.SliceLayout(0, gl.SliceLayout(2, RS)))
+    r2 = gl.arange(0, 8, layout=gl.SliceLayout(0, gl.SliceLayout(1, RS)))
+    return r0[:, None, None] * 16 + r1[None, :, None] * 8 + r2[None, None, :]
+
+
+@gluon.jit
+def row_max(x):
+    return gl.max(gl.max(gl.max(gl.max(x, axis=6), axis=5), axis=4), axis=3)
+
+
+@gluon.jit
+def row_sum(x):
+    return gl.sum(gl.sum(gl.sum(gl.sum(x, axis=6), axis=5), axis=4), axis=3)
+
+
+@gluon.jit
+def bcast_rows(r, L: gl.constexpr):
+    """A [MB, 2, 8] row statistic -> broadcastable over accumulator L (any NB; the row bits are shared)."""
+    RL: gl.constexpr = gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(5, gl.SliceLayout(6, L))))
+    r = gl.convert_layout(r, RL, assert_trivial=True)
+    return gl.expand_dims(gl.expand_dims(gl.expand_dims(gl.expand_dims(r, 3), 4), 5), 6)
+
+
+@gluon.jit
+def stage_k_inner(smem, tile, R: gl.constexpr, KK: gl.constexpr):
+    """tile [R, KK] whose k runs along its columns -> smem [(KK/4) R, 4] laid out [k-step, R, 4]."""
+    smem.store(gl.reshape(gl.permute(gl.reshape(tile, [R, KK // 4, 4]), [1, 0, 2]), [(KK // 4) * R, 4]))
+
+
+@gluon.jit
+def stage_k_outer(smem, tile, KK: gl.constexpr, C: gl.constexpr):
+    """tile [KK, C] whose k runs along its rows -> smem [(KK/4) C, 4] laid out [k-step, C, 4]. Conflict-free
+    stores want each thread to hold 4 consecutive k of a column (a [4, x] size_per_thread load layout)."""
+    smem.store(gl.reshape(gl.permute(gl.reshape(tile, [KK // 4, 4, C]), [0, 2, 1]), [(KK // 4) * C, 4]))
+
+
+@gluon.jit
+def mma_k4(sa, a_row0, sb, b_row0, acc, MB: gl.constexpr, NB: gl.constexpr,
+           LA: gl.constexpr, LB: gl.constexpr, L: gl.constexpr):
+    """acc += A[:, 4j:4j+4] @ B[4j:4j+4, :] for the k-step whose rows start at a_row0 / b_row0 of the staged
+    operands (16 MB rows of sa, 16 NB rows of sb)."""
+    fa = sa.slice(a_row0, 16 * MB).load(LA)                  # (mb, mi3, mi2, ni1, mi0 | mi1, ni0)
+    fa = gl.reshape(fa, [MB, 2, 2, 2, 2, 2, 2])
+    fa = gl.convert_layout(gl.reshape(gl.permute(fa, [0, 1, 2, 5, 4, 3, 6]), [MB, 2, 8, 1, 1, 1, 4]), L,
+                           assert_trivial=True)
+    fb = sb.slice(b_row0, 16 * NB).load(LB)                  # (nb, ni3, mi2, ni1, mi0 | mi1, ni0)
+    fb = gl.reshape(fb, [NB, 2, 2, 2, 2, 2, 2])
+    fb = gl.convert_layout(gl.reshape(gl.permute(fb, [2, 5, 4, 0, 1, 3, 6]), [1, 1, 8, NB, 2, 1, 4]), L,
+                           assert_trivial=True)
+    return gl.inline_asm_elementwise(MMA_ASM, MMA_CONSTRAINTS, [fa, fb, acc], dtype=gl.float32, is_pure=True,
+                                     pack=8)