nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/moe_device_residency.py

@@ -0,0 +1,141 @@
+"""The expert store's residency kept ON THE CARD for a decode token (S3, 2026-10-28).
+
+An expert store serves the routed layers no card holds whole from a cache of device slots. Its host-side residency
+(`core/runtime/expert_residency.ExpertResidency`) reads every call's routing back to the host, decides the misses, and
+copies them — a host round trip per layer per token that also keeps the decode step from being replayed (Qwen3-30B-A3B
+on 4 V100: 8 store layers ~32 ms of a ~80 ms token). For one token the same LRU runs here:
+
+  1. `expert_slot_assign_kernel` — ONE program, deterministic: every routed expert already resident is a hit (its stamp
+     refreshed to the clock); every miss takes the free slot or the least recently stamped slot this call does not read
+     (a slot this call reads carries the current clock, so it is never the oldest), evicting its previous owner from
+     the slot map. Writes each routed expert's slot, and for a miss the expert to copy (-1 for a hit). The clock moves.
+  2. `expert_miss_copy_kernel` — grid (TOP_K, 3 projections, chunks): a miss copies its expert's three projections from
+     page-locked host memory (read through its device-visible address) into its slot; a hit returns at once.
+
+The decode band then reads the experts through the slots' fixed pointer table, indexed by slot (`moe_decode_vec` half
+kernels). Every launch is the same at every token — pointers to persistent tables — so the step replays.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def expert_slot_assign_kernel(
+    ids_ptr,          # [TOP_K] int64 — the routed experts of this token, the router's order
+    slot_of_ptr,      # [L * E] int32 — slot holding (layer, expert), -1 when none
+    owner_ptr,        # [S] int32 — the (layer * E + expert) a slot holds, -1 when free
+    stamp_ptr,        # [S] int64 — the clock at a slot's last use, -1 when free
+    clock_ptr,        # [1] int64
+    slot_out_ptr,     # [TOP_K] int64 — the slot of each routed expert
+    miss_out_ptr,     # [TOP_K] int64 — the expert to copy into that slot, -1 for a hit
+    layer_base,       # layer index * E
+    S,
+    TOP_K: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+):
+    clock = tl.load(clock_ptr)
+    offs = tl.arange(0, BLOCK_S)
+    in_range = offs < S
+    # hits first: their slots carry the current clock before any miss looks for a victim
+    for k in tl.static_range(TOP_K):
+        e = tl.load(ids_ptr + k)
+        s = tl.load(slot_of_ptr + layer_base + e)
+        if s >= 0:
+            tl.store(stamp_ptr + s, clock)
+            tl.store(slot_out_ptr + k, s.to(tl.int64))
+            tl.store(miss_out_ptr + k, -1)
+    tl.debug_barrier()
+    for k in tl.static_range(TOP_K):
+        e = tl.load(ids_ptr + k)
+        s = tl.load(slot_of_ptr + layer_base + e)
+        if s < 0:
+            stamps = tl.load(stamp_ptr + offs, mask=in_range, other=9223372036854775807)
+            victim = tl.argmin(stamps, axis=0)
+            old = tl.load(owner_ptr + victim)
+            if old >= 0:
+                tl.store(slot_of_ptr + old, -1)
+            tl.store(owner_ptr + victim, (layer_base + e).to(tl.int32))
+            tl.store(slot_of_ptr + layer_base + e, victim.to(tl.int32))
+            tl.store(stamp_ptr + victim, clock)
+            tl.store(slot_out_ptr + k, victim.to(tl.int64))
+            tl.store(miss_out_ptr + k, e)
+        tl.debug_barrier()
+    tl.store(clock_ptr, clock + 1)
+
+
+@triton.jit
+def expert_miss_copy_kernel(
+    miss_ptr,         # [TOP_K] int64 — expert to copy, -1 for a hit
+    slot_ptr,         # [TOP_K] int64 — its slot
+    host_g, host_u, host_d,   # [L * E] int64 — device-visible addresses of the host experts' projections
+    slot_g, slot_u, slot_d,   # [S] int64 — the slots' projection addresses
+    layer_base,
+    n_g, n_u, n_d,    # elements per projection
+    BLOCK: tl.constexpr,
+    W_BF16: tl.constexpr,
+):
+    k = tl.program_id(0).to(tl.int64)
+    p = tl.program_id(1).to(tl.int64)
+    c = tl.program_id(2).to(tl.int64)
+    e_raw = tl.load(miss_ptr + k)
+    # a hit (-1) copies nothing: masked, not returned (a bare return is unstructured control flow some backends do
+    # not lower); its address loads read row 0 of the tables, a valid index, and the copy's mask is all false
+    miss = e_raw >= 0
+    e = tl.where(miss, e_raw, 0)
+    s = tl.where(miss, tl.load(slot_ptr + k), 0)
+    wty = tl.bfloat16 if W_BF16 else tl.float16
+    if p == 0:
+        src_a = tl.load(host_g + layer_base + e)
+        dst_a = tl.load(slot_g + s)
+        n = n_g
+    elif p == 1:
+        src_a = tl.load(host_u + layer_base + e)
+        dst_a = tl.load(slot_u + s)
+        n = n_u
+    else:
+        src_a = tl.load(host_d + layer_base + e)
+        dst_a = tl.load(slot_d + s)
+        n = n_d
+    src = tl.cast(src_a, tl.pointer_type(wty), bitcast=True)
+    dst = tl.cast(dst_a, tl.pointer_type(wty), bitcast=True)
+    offs = c * BLOCK + tl.arange(0, BLOCK)
+    m = (offs < n) & miss
+    tl.store(dst + offs, tl.load(src + offs, mask=m), mask=m)
+
+
+def measure_miss_copy(nbytes: int, reps: int, blocks, warps, experts: int) -> dict:
+    """`expert_miss_copy_kernel` on this device, launched as the engine's expert store launches it: `nbytes` in
+    page-locked host memory as `experts` missed experts of three projections each, copied into device slots at every
+    candidate (block, warps), `reps` launches each after one; {"block", "warps", "gbps"} of the fastest, every
+    candidate's rate in "grid". On a V100-SXM2 (2026-11-03): 1024 x 16 at 5.4-5.5 GB/s, where a fixed 4096 x 4 gave
+    3.35."""
+    import time
+    import numpy as np
+    from nbxtensor.memory.allocator import DeviceAllocator as DA
+    from nbxtensor.tensor import NBXTensor
+    from nbxtensor.dtype.nbx_dtype import NBXDtype
+    n = int(nbytes) // (2 * 3 * int(experts))                 # fp16 elements per projection
+    host = [[NBXTensor.empty_cpu((n,), NBXDtype.float16, pinned=True) for _ in range(3)] for _ in range(experts)]
+    slots = [[NBXTensor.empty((n,), dtype=NBXDtype.float16, device="cuda:0") for _ in range(3)] for _ in range(experts)]
+    dev = lambda a: NBXTensor.from_numpy(np.asarray(a, dtype=np.int64)).to("cuda:0")          # noqa: E731
+    host_tabs = [dev([host[e][p].data_ptr() for e in range(experts)]) for p in range(3)]
+    slot_tabs = [dev([slots[e][p].data_ptr() for e in range(experts)]) for p in range(3)]
+    miss = dev(list(range(experts)))
+    moved = experts * 3 * n * 2
+    grid = {}
+    for block in blocks:
+        for w in warps:
+            def launch(block=block, w=w):
+                expert_miss_copy_kernel[(experts, 3, triton.cdiv(n, block))](
+                    miss, miss, *host_tabs, *slot_tabs, 0, n, n, n, BLOCK=block, W_BF16=False, num_warps=w)
+            launch()
+            DA.device_synchronize(0)
+            t0 = time.perf_counter()
+            for _ in range(reps):
+                launch()
+            DA.device_synchronize(0)
+            grid[f"{block}x{w}"] = round(moved * reps / (time.perf_counter() - t0) / 1e9, 2)
+    best = max(grid, key=grid.get)
+    block, w = (int(v) for v in best.split("x"))
+    return {"block": block, "warps": w, "gbps": grid[best], "grid": grid}

nbxtensor/ops/moe_wrappers.py

@@ -0,0 +1,345 @@
+"""Mixture-of-experts wrappers: the fused MoE launches (dense and wna16). Cut out of wrappers.py by family (M2, 2026-10-29), unchanged; shared launch layer and run context stay in wrappers.py."""
+
+from __future__ import annotations
+
+import os
+import triton
+from nbxtensor.ops.autotune_bucket import bucket_of as _bucket_of
+from nbxtensor.ops import launch_keys as _lk
+from nbxtensor.tensor import NBXTensor
+from nbxtensor.dtype.nbx_dtype import NBXDtype, dtype_size
+from nbxtensor.memory.allocator import DeviceAllocator
+from nbxtensor.backend.runtime import set_device
+from nbxtensor.shape import MAX_NDIM as _NBX_MAX_NDIM, broadcast_shapes, max_linear_offset
+from nbxtensor.dtype.itemsize import canonical_name
+from nbxtensor.dtype.promotion import binary_result, float_math_input
+from nbxtensor.dtype.nbx_dtype import parse_dtype
+from nbxtensor.memory.allocator import DeviceOOMError
+from nbxtensor.backend.capabilities import fa_min_tile as _fa_min_tile
+from nbxtensor.backend.capabilities import reduce_tile_max as _reduce_tile_max
+from nbxtensor.backend.faults import device_fault_buffer, device_fault_code_cached, fault_channel
+from nbxtensor.tensor import _upload_int64_array as _nbx_upload_int64, _saturating_cast as _nbx_saturating_cast
+from nbxtensor.ops._configs import sdpa_block_ceiling as _sdpa_block_ceiling
+from nbxtensor.ops._configs import sdpa_launch_meta as _sdpa_launch_meta
+from nbxtensor.ops._configs import matrix_unit as _matrix_unit, matrix_unit_mm_tile as _matrix_unit_mm_tile
+from nbxtensor.ops._configs import largest_tile_within_smem as _largest_tile_within_smem
+from nbxtensor.backend import launcher as _nbx_launcher
+from nbxtensor.backend.launcher import install as _install_launcher
+from nbxtensor.ops.relu import relu_forward_kernel
+from nbxtensor.ops.silu import silu_forward_kernel
+from nbxtensor.ops.gelu import gelu_forward_kernel
+from nbxtensor.ops.sigmoid import sigmoid_forward_kernel
+from nbxtensor.ops.tanh import tanh_forward_kernel
+from nbxtensor.ops.hardsigmoid import hardsigmoid_forward_kernel
+from nbxtensor.ops.hardswish import hardswish_forward_kernel
+from nbxtensor.ops.leaky_relu import leaky_relu_forward_kernel
+from nbxtensor.ops.elu import elu_forward_kernel
+from nbxtensor.ops.mish import mish_forward_kernel
+from nbxtensor.ops.selu import selu_forward_kernel
+from nbxtensor.ops.neg import neg_forward_kernel
+from nbxtensor.ops.exp import exp_forward_kernel
+from nbxtensor.ops.sin import sin_forward_kernel
+from nbxtensor.ops.cos import cos_forward_kernel
+from nbxtensor.ops.rsqrt import rsqrt_forward_kernel
+from nbxtensor.ops.sqrt import sqrt_forward_kernel
+from nbxtensor.ops.abs import abs_forward_kernel
+from nbxtensor.ops.log import log_forward_kernel
+from nbxtensor.ops.reciprocal import reciprocal_forward_kernel
+from nbxtensor.ops.pow import pow_forward_kernel
+from nbxtensor.ops.clamp import clamp_forward_kernel
+from nbxtensor.ops.erf import erf_forward_kernel
+from nbxtensor.ops.copy import copy_forward_kernel
+from nbxtensor.ops.add import add_forward_kernel, add_scalar_kernel, add_scalar_dev_kernel, add_bias_broadcast_kernel
+from nbxtensor.ops.binary_strided import add_strided_nd_kernel, mul_strided_nd_kernel, sub_strided_nd_kernel, div_strided_nd_kernel, add_scalar_strided_nd_kernel, mul_scalar_strided_nd_kernel
+from nbxtensor.ops.mul import mul_forward_kernel, mul_scalar_kernel, mul_scalar_dev_kernel
+from nbxtensor.ops.div import div_forward_kernel, div_scalar_kernel, div_scalar_dev_kernel, int_div_kernel
+from nbxtensor.ops.sub import sub_forward_kernel, rsub_forward_kernel
+from nbxtensor.ops.where import where_forward_kernel
+from nbxtensor.ops.maximum import maximum_forward_kernel
+from nbxtensor.ops.minimum import minimum_forward_kernel
+from nbxtensor.ops.masked_fill import masked_fill_forward_kernel
+from nbxtensor.ops.gt import gt_forward_kernel, gt_scalar_kernel
+from nbxtensor.ops.ge import ge_forward_kernel, ge_scalar_kernel
+from nbxtensor.ops.lt import lt_forward_kernel, lt_scalar_kernel
+from nbxtensor.ops.eq import eq_forward_kernel, eq_scalar_kernel
+from nbxtensor.ops.ne import ne_forward_kernel, ne_scalar_kernel
+from nbxtensor.ops.layernorm import layer_norm_forward_kernel
+from nbxtensor.ops.rmsnorm import rms_norm_forward_kernel
+from nbxtensor.ops.softmax import softmax_forward_kernel, softmax_forward_looped_kernel
+from nbxtensor.ops.matmul import matmul_kernel, addmm_kernel
+from nbxtensor.ops._autotune_policy import NBX_FORCE_FP32_ACCUM
+from nbxtensor.ops.fused_moe import fused_moe_kernel, silu_and_mul_kernel, silu_mul_split_kernel
+from nbxtensor.ops.dtype_convert import bf16_to_fp16_kernel
+from nbxtensor.ops.embedding import embedding_kernel, EMBEDDING_OOB
+from nbxtensor.ops.sum import sum_forward_kernel
+from nbxtensor.ops.mean import mean_forward_kernel
+from nbxtensor.ops.amax import amax_forward_kernel
+from nbxtensor.ops.glu import glu_forward_kernel
+from nbxtensor.ops.le import le_forward_kernel, le_scalar_kernel
+from nbxtensor.ops.softplus import softplus_forward_kernel
+from nbxtensor.ops.dropout import dropout_inference_kernel
+from nbxtensor.ops.upsample_nearest2d import upsample_nearest2d_kernel
+from nbxtensor.ops.groupnorm import group_norm_forward_kernel
+from nbxtensor.ops.index_select import index_select_kernel, INDEX_SELECT_OOB
+from nbxtensor.ops.triu import triu_kernel, triu_batch_kernel
+from nbxtensor.ops.tril import tril_kernel, tril_batch_kernel
+from nbxtensor.ops.argmax import argmax_kernel_1, argmax_kernel_2, argmax_kernel_inner
+from nbxtensor.ops.argmin import argmin_kernel_1, argmin_kernel_2, argmin_kernel
+from nbxtensor.ops.conv2d import conv2d_forward_kernel
+from nbxtensor.ops.depthwise_conv2d import depthwise_conv2d_kernel
+from nbxtensor.ops.batch_norm import batch_norm_forward_kernel
+from nbxtensor.ops.cumsum import scan_part_sum_kernel, add_base_sum_kernel
+from nbxtensor.ops.scatter_op import scatter_kernel, scatter_add_kernel, scatter_reduce_amax_kernel, scatter_reduce_amin_kernel
+from nbxtensor.ops.gather_op import gather_kernel
+from nbxtensor.ops.avg_pool2d import avg_pool2d_forward_kernel
+from nbxtensor.ops.max_pool2d import max_pool2d_forward_kernel
+from nbxtensor.ops.max_pool3d import max_pool3d_forward_kernel
+from nbxtensor.ops.cross_entropy import cross_entropy_loss_forward_kernel
+from nbxtensor.ops.exp2 import exp2_forward_kernel
+from nbxtensor.ops.tan import tan_forward_kernel
+from nbxtensor.ops.celu import celu_forward_kernel
+from nbxtensor.ops.log_sigmoid import log_sigmoid_forward_kernel
+from nbxtensor.ops.isfinite import isfinite_forward_kernel
+from nbxtensor.ops.isinf import isinf_forward_kernel
+from nbxtensor.ops.isnan import isnan_forward_kernel
+from nbxtensor.ops.nan_to_num import nan_to_num_forward_kernel
+from nbxtensor.ops.threshold import threshold_forward_kernel
+from nbxtensor.ops.floor import floor_forward_kernel, ceil_forward_kernel, round_forward_kernel, trunc_forward_kernel
+from nbxtensor.ops.addcdiv import addcdiv_forward_kernel
+from nbxtensor.ops.addcmul import addcmul_forward_kernel
+from nbxtensor.ops.all_reduce import all_kernel_mid, all_kernel_result, all_kernel_dim
+from nbxtensor.ops.any_reduce import any_kernel_mid, any_kernel_result, any_kernel_dim
+from nbxtensor.ops.bitwise_and import bitwise_and_forward_kernel
+from nbxtensor.ops.bitwise_or import bitwise_or_forward_kernel
+from nbxtensor.ops.bitwise_not import bitwise_not_forward_kernel
+from nbxtensor.ops.logical_and import logical_and_forward_kernel
+from nbxtensor.ops.logical_or import logical_or_forward_kernel
+from nbxtensor.ops.logical_not import logical_not_forward_kernel
+from nbxtensor.ops.lerp import lerp_tensor_forward_kernel, lerp_scalar_head_kernel, lerp_scalar_tail_kernel
+from nbxtensor.ops.dot_op import dot_kernel_small, dot_kernel_partial, dot_kernel_reduce
+from nbxtensor.ops.flip_op import flip_1d_kernel, flip_strided_kernel
+from nbxtensor.ops.cat_op import cat_copy_kernel_4
+from nbxtensor.ops.fill_op import fill_kernel
+from nbxtensor.ops.arange_op import arange_kernel
+from nbxtensor.ops.prod import prod_kernel_mid, prod_kernel_result, prod_kernel
+from nbxtensor.ops.min_reduce import min_kernel_mid, min_kernel_result, min_kernel
+from nbxtensor.ops.max_reduce import max_kernel_mid, max_kernel_result, max_kernel
+from nbxtensor.ops.baddbmm_op import baddbmm_kernel
+from nbxtensor.ops.addmv_op import addmv_kernel
+from nbxtensor.ops.mse_loss import mse_loss_partial_kernel, mse_loss_reduce_kernel, mse_loss_none_kernel
+from nbxtensor.ops.nllloss import nll_loss_forward_kernel
+from nbxtensor.ops.std import std_map_kernel, std_reduce_kernel, std_dim_kernel
+from nbxtensor.ops.var import var_kernel_1, var_kernel_2, var_welford_kernel
+from nbxtensor.ops.index_add import index_add_gather_kernel
+from nbxtensor.ops.index_put_op import index_put_kernel, INDEX_PUT_OOB
+from nbxtensor.ops.sort_op import (radix_sort_histogram_kernel,
+                          radix_sort_scatter_kernel,
+                          radix_sort_tile_counts_kernel,
+                          radix_sort_tile_prefix_kernel)
+from nbxtensor.ops.rope import rope_forward_kernel
+from nbxtensor.ops.pixel_shuffle import (
+    pixel_shuffle_kernel,
+    pixel_shuffle_broadcast_aware_kernel,
+)
+from nbxtensor.ops.pixel_unshuffle import pixel_unshuffle_kernel
+from nbxtensor.ops.upsample_bilinear2d import upsample_bilinear2d_kernel
+from nbxtensor.ops.adaptive_avg_pool2d import adaptive_avg_pool2d_kernel
+from nbxtensor.ops.conv_depthwise2d import conv_depthwise2d_kernel
+from nbxtensor.ops.scaled_softmax import scaled_softmax_kernel
+from nbxtensor.ops.vector_norm import (
+    l2_norm_kernel, l1_norm_kernel, linf_norm_kernel, l0_norm_kernel, lp_norm_kernel,
+    l2_norm_pass1_kernel, l2_norm_pass2_kernel,
+    linf_norm_pass1_kernel, linf_norm_pass2_kernel,
+)
+from nbxtensor.ops.addr import addr_kernel
+from nbxtensor.ops.logical_xor import logical_xor_kernel
+from nbxtensor.ops.bitwise_left_shift import bitwise_left_shift_kernel
+from nbxtensor.ops.bitwise_right_shift import bitwise_right_shift_kernel
+from nbxtensor.ops.rand_op import rand_kernel, randn_kernel
+from nbxtensor.ops.clamp_min import clamp_min_forward_kernel
+from nbxtensor.ops.conv_transpose2d import conv_transpose2d_kernel
+from nbxtensor.ops.cummax import scan_part_max_kernel
+from nbxtensor.ops.cummin import scan_part_min_kernel
+from nbxtensor.ops.count_nonzero import count_nonzero_kernel
+from nbxtensor.ops.trace_op import trace_kernel
+from nbxtensor.ops.linspace_op import linspace_kernel
+from nbxtensor.ops.full_op import full_kernel
+from nbxtensor.ops.stack_op import stack_copy_kernel
+from nbxtensor.ops.repeat_op import repeat_1d_kernel
+from nbxtensor.ops.upsample_nearest3d import upsample_nearest3d_kernel
+from nbxtensor.ops.weight_norm import weight_norm_kernel_first, weight_norm_kernel_last
+from nbxtensor.ops.repeat_interleave import repeat_interleave_tensor_kernel
+from nbxtensor.ops.elementwise_launch import ew_meta as _ew_meta, ew_grid as _ew_grid
+from nbxtensor.ops import tiling_sizes as _ts
+import math
+
+
+
+# ===========================================================================
+# FUSED MOE WRAPPER
+# ===========================================================================
+
+# Block sizes for fused MoE kernel (tuned for decode batch=1)
+_MOE_BM = 16
+_MOE_BN = 64
+_MOE_BK = 32
+_MOE_GROUP = 8
+
+
+def invoke_fused_moe(
+    hidden_states,            # NBXTensor [M, K] — activations
+    expert_ptrs,              # NBXTensor [E] int64 — absolute expert weight pointers
+    output,                   # NBXTensor [M * top_k, N] — pre-allocated output
+    topk_weights,             # NBXTensor [M * top_k] — flat routing scores
+    sorted_token_ids,         # NBXTensor [num_tokens_post_padded] — sorted indices
+    expert_ids,               # NBXTensor [num_blocks] — expert per BLOCK_M group
+    num_tokens_post_padded,   # NBXTensor [1] — total sorted entries
+    N, K,                     # output dim, reduction dim
+    stride_bk, stride_bn,    # weight strides (shared by all experts)
+    top_k,                    # int
+    mul_routed_weight=False,
+    topk_divide=True,
+):
+    """Launch fused MoE grouped GEMM kernel (absolute pointer table).
+
+    One kernel launch handles ALL experts. Each expert's weight pointer is
+    stored as an absolute int64 in the table — no offset arithmetic.
+    """
+    M = hidden_states.shape[0]
+    # num_valid_tokens MUST be the true assignment count (tokens × top_k)
+    # — topk_weights is [M_tokens * top_k] on every pass. Deriving it
+    # from hidden_states.shape[0] was correct for the gate/up passes
+    # (input [M_tokens, K]) but WRONG for the down pass (input
+    # `activated` is [padded, N]): the inflated bound turned the
+    # sentinel lanes' token_mask TRUE, and the moe_weight load
+    # (mul_routed_weight, down pass only) read past the [n] routing
+    # buffer — a layout-dependent Warp Illegal Address, named by the
+    # CUDA coredump at fused_moe.py:107 on the Qwen3-Omni audio leg
+    # (immune to every sync barrier, hidden under CUDA_LAUNCH_BLOCKING
+    # because the heap layout — not timing — decided whether the OOB
+    # page was mapped).
+    grid, args, kwargs = _fused_moe_launch(hidden_states, expert_ptrs, output, topk_weights, sorted_token_ids,
+                                           expert_ids, num_tokens_post_padded, N, K, stride_bk, stride_bn, top_k,
+                                           mul_routed_weight, topk_divide)
+    set_device(hidden_states)
+    fused_moe_kernel[grid](*args, **kwargs)
+
+
+def _fused_moe_launch(hidden_states, expert_ptrs, output, topk_weights, sorted_token_ids, expert_ids,
+                      num_tokens_post_padded, N, K, stride_bk, stride_bn, top_k, mul_routed_weight, topk_divide):
+    """(grid, args, kwargs) of `fused_moe_kernel` for this launch — one definition for the launch and for the key its
+    route reads (`fused_moe_unit_tile`)."""
+    M = hidden_states.shape[0]
+    num_valid_tokens = topk_weights.shape[0]
+    EM = sorted_token_ids.shape[0]
+    if M < _MOE_BM:
+        EM = min(EM, M * top_k * _MOE_BM)
+    # the tile beyond the alignment's row block is the key's certified configuration (`fused_moe._FUSED_MOE_CONFIGS`)
+    grid = lambda META: (triton.cdiv(EM, _MOE_BM) * triton.cdiv(N, META['BLOCK_SIZE_N']),)    # noqa: E731
+    args = (hidden_states, expert_ptrs, output, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded,
+            N, K, EM, num_valid_tokens, _lk.bucket_of("M", num_valid_tokens), expert_ptrs.numel(),
+            hidden_states.stride(0), hidden_states.stride(1), stride_bk, stride_bn, output.stride(0), output.stride(1))
+    kwargs = dict(BLOCK_SIZE_M=_MOE_BM, MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=top_k,
+                  compute_type=hidden_states.dtype, TOPK_DIVIDE=topk_divide)
+    return grid, args, kwargs
+
+
+def fused_moe_unit_tile(hidden_states, expert_ptrs, output, topk_weights, sorted_token_ids, expert_ids,
+                        num_tokens_post_padded, N, K, stride_bk, stride_bn, top_k, mul_routed_weight=False,
+                        topk_divide=True):
+    """The matrix unit's tile (a `matrix_unit.mm` row) a fused MoE launch runs at when its certificate measured the
+    unit faster than the FMA kernel (`autotune_certified.unit_choice`, the FMA launch's key), else None. The unit takes
+    experts of its operand dtype only (no split), and a launch with no certificate keeps the FMA kernel; a census
+    shadow and the certifier's sweep take the FMA kernel (the unit's tiles are measured beside it,
+    `matrix_unit_runs`). The arguments are `invoke_fused_moe`'s at the FMA alignment."""
+    from nbxtensor.ops._configs import matrix_unit
+    from nbxtensor.ops import autotune_certified as _cert
+    from nbxtensor.ops.launch_observer import launch_observer
+    if hidden_states.nbx_dtype != NBXDtype.float16 or not (matrix_unit() or {}).get("mm"):
+        return None
+    if launch_observer().active() or _cert.certifying():
+        return None
+    _grid, args, kwargs = _fused_moe_launch(hidden_states, expert_ptrs, output, topk_weights, sorted_token_ids,
+                                            expert_ids, num_tokens_post_padded, N, K, stride_bk, stride_bn, top_k,
+                                            mul_routed_weight, topk_divide)
+    chosen = _cert.unit_choice(fused_moe_kernel, args, kwargs)
+    return chosen[1] if chosen is not None and chosen[0] == "unit" else None
+
+
+def fused_moe_unit_tiles(hidden_states):
+    """The matrix-unit tiles a fused MoE launch on these activations can run at (`matrix_unit.mm` of the profile in
+    force, the unit's operand dtype only): the candidates its certification measures beside the FMA kernel."""
+    from nbxtensor.ops._configs import matrix_unit
+    from nbxtensor.ops.autotune_certified import _UNIT_TILE_FIELDS
+    if hidden_states.nbx_dtype != NBXDtype.float16:
+        return []
+    return [{f: int(row[f]) for f in _UNIT_TILE_FIELDS} for row in ((matrix_unit() or {}).get("mm") or [])]
+
+
+def invoke_fused_moe_m8n8k4(hidden_states, expert_ptrs, output, topk_weights, sorted_token_ids, expert_ids,
+                            num_tokens_post_padded, N, K, stride_bk, stride_bn, top_k, tile, mul_routed_weight=False,
+                            topk_divide=True):
+    """`invoke_fused_moe` on the m8n8k4 matrix unit (`fused_moe_m8n8k4_kernel`) at `tile`; the routing aligned at the
+    tile's rows (`moe_align_block_size(..., tile["block_m"], ...)`)."""
+    from nbxtensor.ops.fused_moe_m8n8k4 import fused_moe_m8n8k4_kernel
+    from nbxtensor.ops.matmul_m8n8k4 import mm_layouts
+    BM, BN, BK = int(tile["block_m"]), int(tile["block_n"]), int(tile["block_k"])
+    WM, WN = int(tile["warps_m"]), int(tile["warps_n"])
+    EM = sorted_token_ids.shape[0]
+    b_k_contig = stride_bk == 1 and N > 1
+    set_device(hidden_states)
+    fused_moe_m8n8k4_kernel[(triton.cdiv(EM, BM) * triton.cdiv(N, BN),)](
+        hidden_states, expert_ptrs, output, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded,
+        N, K, EM, topk_weights.shape[0],
+        hidden_states.stride(0), hidden_states.stride(1), stride_bk, stride_bn, output.stride(0), output.stride(1),
+        MUL_ROUTED_WEIGHT=bool(mul_routed_weight), TOP_K=int(top_k), TOPK_DIVIDE=bool(topk_divide),
+        B_K_CONTIG=b_k_contig, MB=BM // 16, NB=BN // 16, BK=BK, GROUP_M=int(tile["group_m"]),
+        num_warps=WM * WN, **mm_layouts(BM // 16, BN // 16, WM, WN, BK, b_k_contig))
+
+
+def invoke_fused_moe_wna16(
+    hidden_states,            # NBXTensor [M, K] — activations
+    qw_ptrs, sc_ptrs, mn_ptrs,  # NBXTensor [E] int64 — triplet tables
+    output,
+    topk_weights,
+    sorted_token_ids,
+    expert_ids,
+    num_tokens_post_padded,
+    N, K,
+    stride_qk, stride_qn,     # packed int32 strides
+    stride_sg, stride_sn,     # scales/mins strides
+    top_k,
+    mul_routed_weight=False,
+    topk_divide=True,
+):
+    """Launch the W4 fused MoE grouped GEMM (int4-g128-asym experts,
+    in-register dequant — byte-gated vs fused_moe_fp32b_kernel).
+    Same launch contract as invoke_fused_moe; _MOE_BK (32) divides the
+    quant group (128) so a K-tile never crosses a group boundary."""
+    from nbxtensor.ops.fused_moe import fused_moe_wna16_kernel
+    M = hidden_states.shape[0]
+    num_valid_tokens = topk_weights.shape[0]
+    EM = sorted_token_ids.shape[0]
+    if M < _MOE_BM:
+        EM = min(EM, M * top_k * _MOE_BM)
+    grid = (triton.cdiv(EM, _MOE_BM) * triton.cdiv(N, _MOE_BN),)
+    set_device(hidden_states)
+
+    fused_moe_wna16_kernel[grid](
+        hidden_states, qw_ptrs, sc_ptrs, mn_ptrs, output,
+        topk_weights,
+        sorted_token_ids, expert_ids, num_tokens_post_padded,
+        N, K, EM, num_valid_tokens,
+        hidden_states.stride(0), hidden_states.stride(1),
+        stride_qk, stride_qn,
+        stride_sg, stride_sn,
+        output.stride(0), output.stride(1),
+        BLOCK_SIZE_M=_MOE_BM, BLOCK_SIZE_N=_MOE_BN, BLOCK_SIZE_K=_MOE_BK,
+        GROUP_SIZE_M=_MOE_GROUP,
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=hidden_states.dtype,
+        QGROUP=128, QPACK=8,
+        TOPK_DIVIDE=topk_divide,
+        num_warps=4, num_stages=2,
+    )

nbxtensor/ops/mse_loss.py

@@ -0,0 +1,79 @@
+"""MSE Loss kernels — extracted from FlagGems.
+
+Pure Triton. ZERO PyTorch imports.
+
+Two-pass reduction:
+  kernel_1: Per-block partial (input - target)^2, reduced by mean or sum.
+  kernel_2: Final reduction across blocks.
+
+Reduction modes: 1 = mean, 2 = sum.
+For reduction=0 (none), use a simple pointwise kernel instead.
+"""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def mse_loss_partial_kernel(
+    inp,
+    target,
+    mid,
+    M,
+    BLOCK_SIZE: tl.constexpr,
+    reduction: tl.constexpr,
+):
+    """Pass 1: per-block squared difference, reduced to partial sums.
+
+    For reduction==1 (mean): each block stores sum((inp-target)^2) / M.
+    For reduction==2 (sum):  each block stores sum((inp-target)^2).
+    """
+    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 < M
+
+    inp_val = tl.load(inp + offset, mask=mask, other=0).to(tl.float32)
+    target_val = tl.load(target + offset, mask=mask, other=0).to(tl.float32)
+    sub = inp_val - target_val
+    pow_val = sub * sub
+
+    # reduction == 1: mean, reduction == 2: sum
+    if reduction == 1:
+        sum_val = tl.sum(pow_val) / M
+    else:
+        sum_val = tl.sum(pow_val)
+
+    tl.store(mid + pid, sum_val)
+
+@triton.jit
+def mse_loss_reduce_kernel(
+    mid,
+    out,
+    mid_size,
+    BLOCK_MID: tl.constexpr,
+):
+    """Pass 2: sum all partial results into a single scalar."""
+    offset = tl.arange(0, BLOCK_MID)
+    mask = offset < mid_size
+    mid_val = tl.load(mid + offset, mask=mask, other=0).to(tl.float32)
+    sum_val = tl.sum(mid_val)
+    tl.store(out, sum_val)
+
+@triton.jit
+def mse_loss_none_kernel(
+    inp_ptr,
+    target_ptr,
+    output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Pointwise (no reduction): output[i] = (input[i] - target[i])^2."""
+    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(inp_ptr + offset, mask=mask, other=0).to(tl.float32)
+    y = tl.load(target_ptr + offset, mask=mask, other=0).to(tl.float32)
+
+    diff = x - y
+    result = diff * diff
+    tl.store(output_ptr + offset, result, mask=mask)

nbxtensor/ops/mul.py

@@ -0,0 +1,51 @@
+"""Multiply — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def mul_forward_kernel(
+    x_ptr, y_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x * y (tensor * tensor)"""
+    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 * y, mask=mask)
+
+@triton.jit
+def mul_scalar_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x * scalar"""
+    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, x * scalar, mask=mask)
+
+
+@triton.jit
+def mul_scalar_dev_kernel(
+    x_ptr, s_ptr, output_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """out = x * s where s is a 0-d DEVICE tensor. Bit-exact mirror
+    of mul_scalar_kernel (host float(s) -> f32 arg == load -> f64 ->
+    f32). Device-scalar increment 2026-08-15."""
+    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
+    s = tl.load(s_ptr).to(tl.float64).to(tl.float32)
+    x = tl.load(x_ptr + offset, mask=mask)
+    tl.store(output_ptr + offset, x * s, mask=mask)

nbxtensor/ops/nan_to_num.py

@@ -0,0 +1,23 @@
+"""NanToNum — pure @triton.jit kernel. Logic from FlagGems."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def nan_to_num_forward_kernel(
+    input_ptr, output_ptr,
+    n_elements,
+    nan_val, posinf_val, neginf_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)
+    pos_inf = float('inf')
+    neg_inf = float('-inf')
+    out = tl.where(x != x, nan_val,
+          tl.where(x == pos_inf, posinf_val,
+          tl.where(x == neg_inf, neginf_val, x)))
+    tl.store(output_ptr + offset, out, mask=mask)

nbxtensor/ops/ne.py

@@ -0,0 +1,32 @@
+"""Not equal — pure @triton.jit kernel."""
+
+import triton
+import triton.language as tl
+
+@triton.jit
+def ne_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, (x != y).to(tl.int1), mask=mask)
+
+@triton.jit
+def ne_scalar_kernel(
+    x_ptr, output_ptr,
+    n_elements,
+    scalar,
+    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)
+    tl.store(output_ptr + offset, (x != scalar).to(tl.int1), mask=mask)