nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/launch_keys.py

@@ -0,0 +1,1054 @@
+"""The autotuned kernels' launch decisions as PURE functions of shapes and dtypes.
+
+The owner's ruling (2026-09-29 01:37): the census is a DERIVATION from the graphs, never an
+execution. Every key a run forms is a function of the op's operand shapes, operand dtypes and
+attributes, the hardware's `has_native_bf16`, and the profile's bucket ladders — so the key is
+computed here, from those alone, and the wrappers make their launch decisions through the SAME
+functions: one implementation, which the derived census calls without launching anything and the
+wrappers call before they launch. A census shadow run is the proof that the two agree
+(`census.expect_launches`: every recorded key must be one the op's derivation named).
+
+Nothing here touches a tensor or a device. The dtypes are `NBXDtype` members; a key's dtype tags
+are the launcher's own spelling (`str` of the Triton dtype, via `_get_tl_dtype`), in the kernel's
+argument order, exactly as `autotune_cache.key_of` appends them.
+"""
+from __future__ import annotations
+
+from typing import List, Optional, Tuple
+
+from nbxtensor.ops import tiling_sizes as _ts
+from nbxtensor.ops.autotune_bucket import bucket_of
+from nbxtensor.dtype.nbx_dtype import NBXDtype
+from nbxtensor.tensor import _get_tl_dtype
+
+MATMUL = "nbxtensor.ops.matmul.matmul_kernel"
+ADDMM = "nbxtensor.ops.matmul.addmm_kernel"
+BADDBMM = "nbxtensor.ops.baddbmm_op.baddbmm_kernel"
+CONV2D = "nbxtensor.ops.conv2d.conv2d_forward_kernel"
+DEPTHWISE = "nbxtensor.ops.depthwise_conv2d.depthwise_conv2d_kernel"
+SOFTMAX = "nbxtensor.ops.softmax.softmax_forward_kernel"
+SOFTMAX_LOOPED = "nbxtensor.ops.softmax.softmax_forward_looped_kernel"
+ELEMENTWISE = "nbxtensor.ops.elementwise_launch.ew_launch_kernel"
+GROUP_NORM = "nbxtensor.ops.groupnorm.group_norm_forward_kernel"
+ROW_REDUCE = {"mean": "nbxtensor.ops.mean.mean_forward_kernel",
+              "sum": "nbxtensor.ops.sum.sum_forward_kernel",
+              "amax": "nbxtensor.ops.amax.amax_forward_kernel"}
+UNIT_FLASH = "nbxtensor.ops.flash_attention_m8n8k4.flash_attention_m8n8k4_kernel"
+FLASH = "nbxtensor.ops.flash_attention.flash_attention_certified_kernel"
+LAYER_NORM = "nbxtensor.ops.layernorm.layer_norm_forward_kernel"
+RMS_NORM = "nbxtensor.ops.rmsnorm.rms_norm_forward_kernel"
+GEMV_VEC = "nbxtensor.ops.gemv_vec.gemv_vec_kernel"
+ADDMV = "nbxtensor.ops.addmv_op.addmv_kernel"
+FUSED_MOE = "nbxtensor.ops.fused_moe.fused_moe_kernel"
+MOE_GATEUP_HALF = "nbxtensor.ops.moe_decode_vec.moe_gateup_vec_half_kernel"
+MOE_DOWN_HALF = "nbxtensor.ops.moe_decode_vec.moe_down_split_vec_half_kernel"
+MOE_PART_REDUCE = "nbxtensor.ops.moe_decode_vec.moe_part_reduce_kernel"
+DECODE_VEC = "nbxtensor.ops.decode_attn_vec.decode_attn_vec_split_kernel"
+PAGED_DECODE = "nbxtensor.ops.paged_decode_attn.paged_decode_attn_split_kernel"
+DECODE_REDUCE = "nbxtensor.ops.flash_decode.flash_decode_reduce_kernel"
+MOE_ALIGN = tuple(f"nbxtensor.ops.moe_align.moe_align_stage{i}_kernel" for i in (1, 2, 3))
+
+F32, BF16, F16 = NBXDtype.float32, NBXDtype.bfloat16, NBXDtype.float16
+_HALF = (F16, BF16)
+_WIDEST_FIRST = (F32, BF16, F16)
+
+Launch = Tuple[str, tuple]          # (kernel qualified name, key tuple as key_of forms it)
+
+
+def tag(dt: NBXDtype) -> str:
+    """A dtype as the autotune key spells it (`str` of the Triton dtype the tensor reports)."""
+    return str(_get_tl_dtype(dt))
+
+
+def _widest(a: NBXDtype, b: NBXDtype) -> NBXDtype:
+    return next(d for d in _WIDEST_FIRST if d in (a, b))
+
+
+def matmul_out_dtype(dt: NBXDtype, M: int, force_fp32: bool, native_bf16: bool,
+                     force_accum: bool = False, contract_half: bool = False) -> NBXDtype:
+    """The store dtype of mm/bmm/addmm/mv/addmv (see `wrappers._matmul_out_dtype` for the why):
+    fp32 for any half input under NBX_FORCE_FP32_ACCUM; fp32 for fp16 on hardware without native
+    bf16 UNLESS the component's precision contract holds (`contract_half`: the calibration measured
+    its activations fp16-safe — the compiled branch stores fp16 there, `_FP16_GEMM_OPS` under the
+    contract, and a half graph is served half on both engines, the owner 2026-10-25 20:36); fp32
+    for a half input when M <= 4 or when forced, on hardware without native bf16 and no contract;
+    otherwise the input dtype. An island's GEMM reaches here with fp32 operands and stores fp32."""
+    half = dt in _HALF
+    if force_accum and half:
+        return F32
+    # the contract lifts the store for the product rows the calibration observed — a graph op's prefill or spatial
+    # rows; a decode row (M <= 4) keeps the fp32 store that catches decode-loop drift (below): stored fp16 there,
+    # MiniCPM-o-4_5 decoded "The!!!!!" and VibeVoice-1.5B said "Thank you." (gate tha3, 2026-10-26; both correct
+    # with the decode rows conservative, bisected on a card)
+    if dt == F16 and contract_half and M > 4:
+        return F16
+    if dt == F16 and not native_bf16:
+        return F32
+    # ONE precision contract (the owner, 2026-10-25 15:26): on hardware that computes bf16 natively a half input
+    # stores half — the batched attention GEMM (force_fp32) included; fp32 is stored only for an island or where the
+    # hardware has no native half path (above). The compiled branch stores bf16 there (`_wrap_lower_precision`); the
+    # Triton branch stored fp32 and spread it (2 017 policy sites on Sana, the Mac). A DECODE ROW (M <= 4) keeps its
+    # fp32 store on every hardware: stored half on the M4 Pro (native bf16), whisper-large-v3-turbo's Triton decoder
+    # drifted from its first tokens ("Sierrachan... lateral lateral...", stable f8def0b9 -> 58b39d7e; bisected on the
+    # Mac 2026-10-28 to d4809a47, this hunk alone) — the same drift the fp16 decode rows showed (MiniCPM-o "The!!!!!").
+    if half and (M <= 4 or (force_fp32 and not native_bf16)):
+        return F32
+    return dt
+
+
+def mm_dtypes(a: NBXDtype, b: NBXDtype, native_bf16: bool, force_accum: bool = False):
+    """`mm`'s operand decisions: (a at launch, b at launch, PROMOTE_A, PROMOTE_B). An fp16
+    activation on hardware without native bf16 is widened IN the kernel (PROMOTE_A, a stays fp16
+    in memory); a narrow weight under an fp32 activation is widened in the kernel (PROMOTE_B, any
+    hardware); any other mismatch widens both in memory to the widest of the pair."""
+    if force_accum and a in _HALF:
+        a = F32
+    if force_accum and b in _HALF:
+        b = F32
+    promote_a = (not native_bf16) and a == F16
+    a_eff = F32 if promote_a else a
+    promote_b = a_eff == F32 and b in _HALF
+    if a_eff != b and not promote_b:
+        if promote_a:
+            a, promote_a = F32, False
+        w = _widest(a, b)
+        a, b = w, w
+    return a, b, promote_a, promote_b
+
+
+def addmm_dtypes(a: NBXDtype, b: NBXDtype, bias: NBXDtype, native_bf16: bool,
+                 force_accum: bool = False):
+    """`addmm`'s operand decisions: (a, b, bias at launch, PROMOTE_A, PROMOTE_B, PROMOTE_BIAS).
+    As `mm`, except PROMOTE_B is gated on hardware without native bf16 and on an fp16 weight
+    only; the bias is widened in the kernel under an fp32 activation, else cast to a's dtype."""
+    if force_accum and a in _HALF:
+        a = F32
+    if force_accum and b in _HALF:
+        b = F32
+    promote_a = (not native_bf16) and a == F16
+    a_eff = F32 if promote_a else a
+    promote_b = (not native_bf16) and a_eff == F32 and b == F16
+    if a_eff != b and not promote_b:
+        if promote_a:
+            a, promote_a = F32, False
+        w = _widest(a, b)
+        a, b = w, w
+        a_eff = a
+    promote_bias = bias != a_eff and a_eff == F32 and bias in _HALF
+    if bias != a_eff and not promote_bias:
+        bias = a_eff
+    return a, b, bias, promote_a, promote_b, promote_bias
+
+
+def bmm_dtypes(a: NBXDtype, b: NBXDtype, native_bf16: bool, force_accum: bool = False):
+    """`bmm`'s operand decisions: (a, b at launch, PROMOTE_B). An fp16 activation on hardware
+    without native bf16 is widened in memory (bmm has no PROMOTE_A); an fp16 weight under fp32 is
+    widened in the kernel on that hardware; any other mismatch widens both to the widest."""
+    if force_accum and a in _HALF:
+        a = F32
+    if force_accum and b in _HALF:
+        b = F32
+    if (not native_bf16) and a == F16:
+        a = F32
+    promote_b = (not native_bf16) and a == F32 and b == F16
+    if a != b and not promote_b:
+        w = _widest(a, b)
+        a, b = w, w
+    return a, b, promote_b
+
+
+def matrix_unit_operands(a: NBXDtype, b: NBXDtype, M: int, unit: Optional[dict] = None):
+    """(rep_a, rep_b): how the profile's matrix unit carries each operand of a GEMM with M output rows, as they are
+    in memory — the DtypeEngine's `matrix_unit_representation` ("native", "split") — or None: either operand has
+    none, or M does not fill the smallest tile the profile lists for that representation (a row band mostly empty
+    leaves the unit idle; the GEMV / batched kernels own those rows, and a row's reduction order then never
+    depends on how many rows share its launch). The wrappers' `_gemm_unit_tile` asks this same function. `unit`: a
+    matrix unit already read (a planner's, for the device it plans); the profile in force's when None."""
+    from nbxtensor.ops._configs import matrix_unit
+    from nbxtensor.dtype.operands import matrix_unit_representation   # the library's operand rules
+    mu = matrix_unit() if unit is None else unit
+    ra, rb = matrix_unit_representation(a.name, mu), matrix_unit_representation(b.name, mu)
+    if not (ra and rb):
+        return None
+    rows = (mu["fp32_split"] if "split" in (ra, rb) else mu)["mm"]
+    return (ra, rb) if M >= min(int(r["block_m"]) for r in rows) else None
+
+
+def matrix_unit_takes(a: NBXDtype, b: NBXDtype, M: int) -> bool:
+    """The profile's matrix unit can run this GEMM (`matrix_unit_operands`). A batched product on it has NO autotune
+    key (`bmm_launches`, `baddbmm_launches` return []); `mm` and `addmm` keep their FMA key, whose certificate chose
+    between the two routes."""
+    return matrix_unit_operands(a, b, M) is not None
+
+
+def gemm_row_representative(m_bucket: int, a: NBXDtype, b: NBXDtype, unit_keyed: bool = False) -> Optional[int]:
+    """The largest row count M whose GEMM launch forms the autotune key of M bucket `m_bucket`: its bucket is
+    `m_bucket` (`bucket_of`) and the matrix unit does not take it (`matrix_unit_takes`). A bucket's top is not always
+    one: on Volta the class 57..64 of an fp32 GEMM keys M_BUCKET 64, and at 64 itself the split matrix unit takes the
+    GEMM (no key) — a certification synthesising M=64 never reached the autotuner (2026-10-19, 172 keys). Every row
+    count down to one is a candidate: a batched GEMM keys its rows however few (`bmm_launches`, no GEMV below five rows
+    as `mm` has) — stopping at five refused every bmm key of buckets 1-4 as a census defect (80 attention keys on the
+    Mac, 2026-10-20). `unit_keyed`: a GEMM whose key the unit's route keeps (mm, addmm: `mm_launches`) — every row
+    count of the bucket forms it. None: no row count forms the key."""
+    m = int(m_bucket)
+    while m >= 1 and bucket_of("M", m) == m_bucket:
+        if unit_keyed or not matrix_unit_takes(a, b, m):
+            return m
+        m -= 1
+    return None
+
+
+def matrix_unit_conv_operands(x: NBXDtype, w: NBXDtype, M: int, N: int, kh: int, kw: int,
+                              unit: Optional[dict] = None):
+    """(the operands' representations, the tile rows) of a convolution's implicit GEMM on the profile's matrix unit
+    (M output positions, N output channels per group, a kh x kw window), or None: the GEMM's own rule
+    (`matrix_unit_operands`) and, for a split operand, a `matrix_unit.fp32_split.conv` row for that window whose tile
+    the GEMM fills (M >= block_m, N >= block_n) — the split was measured faster than the FMA path only there (a 3x3 is
+    slower at every tile: volta.yml). Native operands take the `mm` rows (rows None). The wrappers, their cost and the
+    census ask this one function."""
+    from nbxtensor.ops._configs import matrix_unit
+    reps = matrix_unit_operands(x, w, M, unit)
+    if reps is None or "split" not in reps:
+        return None if reps is None else (reps, None)
+    mu = matrix_unit() if unit is None else unit
+    rows = [r for r in ((mu.get("fp32_split") or {}).get("conv") or [])
+            if [int(v) for v in r["kernel"]] == [int(kh), int(kw)] and M >= int(r["block_m"]) and N >= int(r["block_n"])]
+    return (reps, rows) if rows else None
+
+
+def mv_launches(N: int, K: int, mat: NBXDtype, vec: NBXDtype, native_bf16: bool,
+                force_accum: bool = False, contract_half: bool = False) -> List[Launch]:
+    """`mv_wrapper` on a [N, K] matrix and a [K] vector: the SIMT GEMV's key — the N and K buckets, then the dtype
+    tags of the matrix, the vector and the output (`matmul_out_dtype` of the matrix at one row)."""
+    out = matmul_out_dtype(mat, 1, False, native_bf16, force_accum, contract_half)
+    return [(GEMV_VEC, (bucket_of("N", N), bucket_of("K", K), tag(mat), tag(vec), tag(out)))]
+
+
+def addmv_launches(N: int, K: int, mat: NBXDtype, vec: NBXDtype, bias: NBXDtype, native_bf16: bool,
+                   force_accum: bool = False, contract_half: bool = False) -> List[Launch]:
+    """`addmv_wrapper` on a [N, K] matrix, a [K] vector and a [N] bias: the key — the N and K buckets (K the kernel's
+    `M`), then the dtype tags of the matrix, the vector, the bias and the output (`matmul_out_dtype` at one row)."""
+    out = matmul_out_dtype(mat, 1, False, native_bf16, force_accum, contract_half)
+    return [(ADDMV, (bucket_of("N", N), bucket_of("M", K), tag(mat), tag(vec), tag(bias), tag(out)))]
+
+
+def mm_launches(M: int, K: int, N: int, a: NBXDtype, b: NBXDtype, native_bf16: bool,
+                force_accum: bool = False, contract_half: bool = False) -> List[Launch]:
+    """`mm` (and `mm_epilogue`): M <= 4 runs one GEMV per row on the operands as `mm_dtypes` casts them (the weight
+    transposed to [N, K], each activation row the vector) — one key; otherwise one key, matrix unit or not — row bands
+    above 2^31 output elements keep the whole shape's M bucket."""
+    a, b, promote_a, promote_b = mm_dtypes(a, b, native_bf16, force_accum)
+    if M <= 4:
+        return mv_launches(N, K, b, a, native_bf16, force_accum, contract_half)
+    # where the matrix unit can take it too, the key is the FMA launch's: its certificate measured both routes and
+    # names the faster (`linalg_wrappers._gemm_route`), the unit's tile included
+    out = matmul_out_dtype(a, M, promote_a, native_bf16, force_accum, contract_half)
+    ieee = (not native_bf16) and out == F32
+    return [(MATMUL, (bucket_of("M", M), N, K, ieee, promote_b, tag(a), tag(b), tag(out)))]
+
+
+def addmm_launches(M: int, K: int, N: int, a: NBXDtype, b: NBXDtype, bias: NBXDtype,
+                   native_bf16: bool, force_accum: bool = False, contract_half: bool = False, bias_row: bool = True) -> List[Launch]:
+    """`addmm` on a 2-D activation (an N-D one is flattened to (numel/K, K) first): M <= 4 runs
+    per-row addmv; otherwise one key, whose certificate decides the matrix unit or the FMA kernel where a [N] / [1, N]
+    bias (`bias_row`) lets the unit take it."""
+    a, b, bias, promote_a, promote_b, _ = addmm_dtypes(a, b, bias, native_bf16, force_accum)
+    if M <= 4:
+        # one addmv per row on the operands as cast, the weight transposed to [N, K], the bias its row or the shared one
+        return addmv_launches(N, K, b, a, bias, native_bf16, force_accum, contract_half)
+    out = matmul_out_dtype(a, M, promote_a, native_bf16, force_accum, contract_half)
+    ieee = (not native_bf16) and out == F32
+    return [(ADDMM, (bucket_of("M", M), N, K, ieee, promote_b, tag(a), tag(b), tag(bias), tag(out)))]
+
+
+def bmm_launches(M: int, K: int, N: int, a: NBXDtype, b: NBXDtype, native_bf16: bool,
+                 force_accum: bool = False, contract_half: bool = False) -> List[Launch]:
+    """`bmm`: always one batched launch through `baddbmm_kernel` (HAS_BIAS False, the output
+    passed as the bias pointer), every dim bucketed, the output fp32 for any half input; on the matrix unit
+    (M > 4, operands taken as they are in memory, before the widening) no key."""
+    if matrix_unit_takes(a, b, M):
+        return []
+    a, b, promote_b = bmm_dtypes(a, b, native_bf16, force_accum)
+    out = matmul_out_dtype(a, M, True, native_bf16, force_accum, contract_half)
+    ieee = (not native_bf16) and out == F32
+    return [(BADDBMM, (bucket_of("M", M), bucket_of("N", N), bucket_of("K", K),
+                       ieee, promote_b, False, tag(a), tag(b), tag(out), tag(out)))]
+
+
+def baddbmm_launches(M: int, K: int, N: int, a: NBXDtype, b: NBXDtype, bias: NBXDtype) -> List[Launch]:
+    """`baddbmm_wrapper`: one batched launch at the operands' own dtypes (no promotion, no
+    widened store: the output is batch1's dtype), IEEE_PRECISION and HAS_BIAS set, PROMOTE_B
+    off, every dim bucketed; the bias pointer keyed at its own dtype. On the matrix unit (M > 4) no key."""
+    if matrix_unit_takes(a, b, M):
+        return []
+    return [(BADDBMM, (bucket_of("M", M), bucket_of("N", N), bucket_of("K", K),
+                       True, False, True, tag(a), tag(b), tag(a), tag(bias)))]
+
+
+def mm_out(M: int, a: NBXDtype, b: NBXDtype, native_bf16: bool, force_accum: bool = False, contract_half: bool = False) -> NBXDtype:
+    """The dtype `mm` stores for these operands (its launch or its per-row GEMV alike)."""
+    a, b, promote_a, _ = mm_dtypes(a, b, native_bf16, force_accum)
+    return matmul_out_dtype(a, M, promote_a, native_bf16, force_accum, contract_half)
+
+
+def bmm_out(M: int, a: NBXDtype, b: NBXDtype, native_bf16: bool, force_accum: bool = False, contract_half: bool = False) -> NBXDtype:
+    """The dtype `bmm` stores for these operands: on the matrix unit the operands are taken as they are in memory
+    (no widening, `wrappers.bmm`'s unit branch), elsewhere as `bmm_dtypes` widens them — under the fp16 contract the
+    two differ (fp16 on the unit, fp32 off it on hardware without native bf16)."""
+    if matrix_unit_takes(a, b, M):
+        return matmul_out_dtype(a, M, True, native_bf16, force_accum, contract_half)
+    a, b, _ = bmm_dtypes(a, b, native_bf16, force_accum)
+    return matmul_out_dtype(a, M, True, native_bf16, force_accum, contract_half)
+
+
+def row_breakpoints(lo: int, hi: int) -> List[int]:
+    """The row counts M in (lo, hi] at which a GEMM's launch CLASS changes on the bound profile: the GEMV bound
+    (M <= 4 launches no autotuned kernel), the M bucket (`bucket_of("M", M)`: exact below the matrix unit's smallest
+    tile, bucketed above), and the matrix unit's take-over for the native and the split operand representations
+    (no key from there on). A census that walks a token axis (a prompt's length) derives at every breakpoint and
+    bisects between two: the key set at M is a step function of M only between breakpoints — below and above the
+    keyed band it is empty at both ends, which a bisection alone takes for one class (TinyLlama, 2026-10-19:
+    a 31-token prompt's matmul key missed under certified-only). One source: the launcher's own functions."""
+    lo, hi = int(lo), int(hi)
+    if hi <= lo:
+        return []
+    half, single = NBXDtype.float16, NBXDtype.float32
+    def klass(m: int):
+        return (m <= 4, bucket_of("M", m), matrix_unit_takes(half, half, m), matrix_unit_takes(single, half, m))
+    out: List[int] = []
+    prev = klass(max(lo, 1))
+    for m in range(max(lo, 1) + 1, hi + 1):
+        k = klass(m)
+        if k != prev:
+            out.append(m)
+            prev = k
+    return out
+
+
+def matmul_route(a_rank: int, b_rank: int) -> str:
+    """`matmul_wrapper`'s route by operand ranks (torch.matmul semantics): "mm" (2-D x 2-D), "bmm"
+    (3-D x 3-D), "mv" (2-D x 1-D, the SIMT GEMV), "batched" (N-D x 2-D:
+    the leading dims folded into bmm's batch), "general" (N-D x N-D: both batches collapsed)."""
+    if a_rank == 2 and b_rank == 2:
+        return "mm"
+    if a_rank == 3 and b_rank == 3:
+        return "bmm"
+    if a_rank == 2 and b_rank == 1:
+        return "mv"
+    if a_rank >= 3 and b_rank == 2:
+        return "batched"
+    if a_rank >= 3 and b_rank >= 3:
+        return "general"
+    raise ValueError(f"matmul: unsupported ranks {a_rank} x {b_rank}")
+
+
+def matmul_launches(a_shape, b_shape, a: NBXDtype, b: NBXDtype, native_bf16: bool,
+                    force_accum: bool = False, contract_half: bool = False) -> List[Launch]:
+    """`matmul_wrapper` (and `linear_wrapper`, which calls it with the transposed weight): the
+    launch of the route `matmul_route` names — bmm's key carries no batch."""
+    route = matmul_route(len(a_shape), len(b_shape))
+    if route == "mm":
+        return mm_launches(a_shape[0], a_shape[1], b_shape[1], a, b, native_bf16, force_accum, contract_half)
+    if route == "mv":
+        return mv_launches(a_shape[0], a_shape[1], a, b, native_bf16, force_accum, contract_half)
+    return bmm_launches(a_shape[-2], a_shape[-1], b_shape[-1], a, b, native_bf16, force_accum, contract_half)
+
+
+# ---------------------------------------------------------------------------------------------
+# Recurrence and spectra: `lstm_wrapper` and the DFT route of the FFT wrappers, as the matmul
+# launches their generic calls make.
+# ---------------------------------------------------------------------------------------------
+
+def lstm_launches(B: int, T: int, I: int, H: int, num_layers: int, bidirectional: bool,
+                  x: NBXDtype, native_bf16: bool) -> List[Launch]:
+    """`lstm_wrapper` -> `_lstm_run_direction`, per layer and direction: the input projection
+    x[B,T,I] @ W_ih^T (matmul's N-D x 2-D route: `bmm` over B batches of [T, I]), then one
+    h[B,H] @ W_hh^T `mm` per step. The weights and the initial state are cast to the input's
+    dtype `x` (`cdt`, read once); the elementwise gate arithmetic promotes to the wider operand,
+    so the state after a step is the widest of the projection's, the step mm's and `x`, and a
+    later layer reads the concatenation at its FIRST operand's dtype (`NBXTensor.cat`): the
+    forward direction's first step. The state's dtype is iterated to its fixed point (it can
+    only widen), so the step keys are every dtype the state takes."""
+    nd = 2 if bidirectional else 1
+    out: List[Launch] = []
+    layer_in = x
+    for layer in range(int(num_layers)):
+        first_fwd = None
+        for d in range(nd):
+            out += bmm_launches(T, I if layer == 0 else H * nd, 4 * H, layer_in, x, native_bf16)   # never the contract
+            wx = bmm_out(T, layer_in, x, native_bf16)
+            h, seen, first = x, set(), None
+            for _step in range(max(int(T), 0)):
+                if h in seen:
+                    break
+                seen.add(h)
+                out += mm_launches(B, H, 4 * H, h, x, native_bf16)          # the recurrence: never the contract
+                h = _widest(_widest(wx, mm_out(B, h, x, native_bf16)), x)
+                first = h if first is None else first
+            if d == 0:
+                first_fwd = first if first is not None else x
+        layer_in = first_fwd
+    return list(dict.fromkeys(out))
+
+
+def dft_r2c_launches(M: int, N: int, onesided: bool, native_bf16: bool) -> List[Launch]:
+    """`fft_r2c_wrapper` on a length that is not a power of two: `_dft_r2c`, the frames flattened
+    to [M, N] and widened to fp32, times the fp32 cos and -sin matrices (two `mm`). A power of
+    two runs the radix-2 butterfly (no autotuned kernel)."""
+    if N > 1 and (N & (N - 1)) == 0:
+        return []
+    bins = N // 2 + 1 if onesided else N
+    return mm_launches(M, N, bins, F32, F32, native_bf16)
+
+
+def dft_c2r_launches(M: int, bins: int, N: int, native_bf16: bool) -> List[Launch]:
+    """`fft_c2r_wrapper` -> `_dft_c2r`, any length: the complex64 spectrum's real and imaginary
+    parts [M, bins] times the fp32 inverse matrices [bins, N] (two `mm`, one key)."""
+    return mm_launches(M, bins, N, F32, F32, native_bf16)
+
+
+def stft_launches(lead: int, length: int, n_fft: int, hop_length: Optional[int],
+                  win_length: Optional[int], onesided: Optional[bool],
+                  native_bf16: bool) -> List[Launch]:
+    """`stft_wrapper` on a signal of `lead` rows x `length` samples: frames of win_length every
+    hop_length (`unfold`, defaults n_fft // 4 and n_fft as the wrapper's), each zero-padded to
+    n_fft, then `fft_r2c_wrapper` over [lead x frames, n_fft] — `dft_r2c_launches`'s rule (a
+    power-of-two n_fft runs the butterfly: no autotuned key)."""
+    hop = (n_fft // 4) if hop_length in (None, 0) else int(hop_length)
+    win = n_fft if win_length in (None, 0) else int(win_length)
+    frames = (int(length) - win) // hop + 1
+    return dft_r2c_launches(int(lead) * frames, int(n_fft),
+                            True if onesided is None else bool(onesided), native_bf16)
+
+
+def istft_launches(batch: int, bins: int, frames: int, n_fft: int,
+                   native_bf16: bool) -> List[Launch]:
+    """`istft_wrapper` on a [batch, bins, frames] spectrum: `fft_c2r_wrapper` over the frames
+    [batch x frames, bins] to n_fft samples each (`dft_c2r_launches`); the overlap-add and the
+    window envelope launch no autotuned kernel."""
+    return dft_c2r_launches(int(batch) * int(frames), int(bins), int(n_fft), native_bf16)
+
+
+# ---------------------------------------------------------------------------------------------
+# Attention (scaled_dot_product_attention): the route, then the math route's two bmm launches.
+# ---------------------------------------------------------------------------------------------
+
+def sdpa_operand_dtypes(q: NBXDtype, k: NBXDtype, v: NBXDtype,
+                        q_round: Optional[NBXDtype] = None, *, alignment: Optional[str] = None):
+    """The (q, k, v, q_round) dtypes `scaled_dot_product_attention_wrapper` computes attention
+    with: the DtypeEngine's decision (`nbxtensor/dtype/operands.attention_operand_dtypes`) under
+    `alignment` — `precision.attention_operands` of the hardware the attention runs on. The
+    wrapper and the derived census run ON that hardware and leave it None (the active profile);
+    Prism plans for a profile that need not be the process's and passes that profile's answer
+    (`Profile.attention_operands`, read by `triton/dtype.profile_attention_operand_alignment`). All three ask this one function, so none holds a rule of
+    its own."""
+    from nbxtensor.dtype.operands import attention_operand_alignment, attention_operand_dtypes   # the library's operand rules
+    qj = q_round if q_round is not None else q
+    if qj == k == v:
+        return q, k, v, q_round
+    if alignment is None:
+        from nbxtensor.ops._configs import active_vendor_profile
+        alignment = attention_operand_alignment(active_vendor_profile())
+    out = attention_operand_dtypes(q.name, k.name, v.name, alignment,
+                                   q_round.name if q_round is not None else None)
+    return tuple(None if d is None else NBXDtype[d] for d in out)
+
+
+def decode_vec_takes(headdim: int, headdim_v: int, mask_numel: Optional[int], seqlen_k: int) -> bool:
+    """Whether a single-query attention on the math route runs the vector decode kernel
+    (`decode_vec_launches`) — `_try_decode_vec`'s contract: equal head dims, and a mask that is one bias
+    per key (or none). Outside it, the math route's two bmm launches."""
+    if headdim != headdim_v:
+        return False
+    return mask_numel is None or mask_numel == seqlen_k
+
+
+def _pow2(n: int) -> bool:
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def flash_any_head_dim() -> bool:
+    """The arch's profile MEASURED the tl.dot flash kernel correct at every head dim (`flash.any_head_dim_correct`):
+    a head dim that is not a power of two runs in the kernel's power-of-two head-dim block, the columns past it masked
+    to zero on load (`EVEN_HEADDIM` False), so it routes like a power of two — on the device's scores budget — instead
+    of on `sdpa.non_pow2_head_scores_bytes` to the math attention, which writes the scores. An arch that has not
+    measured it keeps the math route."""
+    from nbxtensor.ops._configs import active_vendor_profile
+    return ((active_vendor_profile().get("flash") or {}).get("any_head_dim_correct")) is True
+
+
+def flash_tiles_certified() -> bool:
+    """The tl.dot flash's tile is certified per key (workstream 1) where the profile declares no hand table of tiles
+    (`sdpa_thresholds`); a profile that declares one launches its rows as before."""
+    from nbxtensor.ops._configs import active_vendor_profile
+    return not active_vendor_profile().get("sdpa_thresholds")
+
+
+def flash_launches(Tq: int, Tk: int, D: int, q: NBXDtype, k: NBXDtype, v: NBXDtype,
+                   bias: Optional[NBXDtype] = None) -> List[Launch]:
+    """The tl.dot flash where its tile is certified (`flash_tiles_certified`): its autotune key — the query rows'
+    bucket, the keys' bucket, the head dim — then, as Triton appends them, the dtypes of its seven tensors (Q, K, V, the
+    bias — the mask's, or Q's for the zero / causal bias the wrapper builds —, the output like Q, the fp32 log-sum-exp
+    and its fp32 scratch). No key where the profile's hand table chooses the tile."""
+    if not flash_tiles_certified():
+        return []
+    return [(FLASH, (bucket_of("M", Tq), bucket_of("N", Tk), int(D),
+                     tag(q), tag(k), tag(v), tag(bias if bias is not None else q), tag(q), tag(F32), tag(F32)))]
+
+
+def sdpa_chunk_rows(bound: int, batch: int, nheads: int, Tq: int, Tk: int,
+                    min_chunk_rows: int, max_chunks: int) -> int:
+    """Query rows per chunk that keep each chunk's fp32 scores within `bound`, aligned to the
+    arch's row block (`memory.sdpa_math_min_chunk_rows`); 0 when the shape cannot chunk within
+    the arch's chunk ceiling (`memory.sdpa_math_max_chunks`). The TilingEngine sizes it
+    (`tiling_sizes.sdpa_chunk_rows`)."""
+    return _ts.sdpa_chunk_rows(bound, batch, nheads, Tq, Tk, min_chunk_rows, max_chunks)
+
+
+def flash_headdim_detour(D: int, enabled: bool = True) -> int:
+    """The head dim the flash route runs at: a power of two >= 128 is zero-padded by one
+    (`scaled_dot_product_attention_wrapper`'s detour around the flash kernel's wrong answers at
+    those dims), and the padded call re-enters the wrapper — routed afresh at D + 1 (and V's head
+    dim + 1). `enabled` False is the wrapper's diagnostic `NBX_D128_DETOUR=0`."""
+    if not (enabled and D >= 128 and (D & (D - 1)) == 0):
+        return D
+    # An arch whose profile MEASURED the flash kernel correct at these dims on its stack runs them
+    # as they are: the defect was a code-generation one of an older Triton (P-FLASH-D128-CORRECTNESS,
+    # Volta; Dao-AILab names the class on A100), and on Triton 3.8 the 33-case attention oracle
+    # passes with the detour off on a V100 (2026-10-03). Every arch that has not measured it keeps
+    # the detour.
+    from nbxtensor.ops._configs import active_vendor_profile
+    if ((active_vendor_profile().get("flash") or {}).get("pow2_head_dim_correct")) is True:
+        return D
+    return D + 1
+
+
+CONV_WEIGHT_LAYOUTS = ("KCRS", "KRSC")
+
+
+def conv_weight_layout() -> str:
+    """The order a convolution weight is stored in on this arch: the hardware profile's
+    `conv.weight_layout`, "KCRS" (the container's own) or "KRSC" (channels innermost: the 3x3
+    taps of one output read one contiguous channel run; +46-63 % on the M4 Pro, 2026-10-07). A
+    profile that does not state it is refused by name; with no profile in force (no device) the
+    weights stay in the container's order."""
+    from nbxtensor.ops._configs import active_vendor_profile
+    profile = active_vendor_profile()
+    if not profile:
+        return "KCRS"
+    layout = (profile.get("conv") or {}).get("weight_layout")
+    if layout not in CONV_WEIGHT_LAYOUTS:
+        raise ValueError(f"the hardware profile states conv.weight_layout={layout!r}; every profile "
+                         f"states one of {CONV_WEIGHT_LAYOUTS}")
+    return layout
+
+
+def conv_weight_krsc(w_shape, groups: int, transposed: bool, layout: str) -> bool:
+    """Whether a convolution weight of this shape is stored KRSC under `layout`: a 2-D, not
+    transposed convolution with a spatial extent (a 1x1 weight is the same memory in both
+    orders) that is not depthwise (one channel a group and a group per output: its kernel reads
+    no channel stride)."""
+    if layout != "KRSC" or transposed or len(w_shape) != 4:
+        return False
+    K, Cg, kh, kw = w_shape
+    return kh * kw > 1 and not (Cg == 1 and groups == K)
+
+
+def krsc_conv_weights(dag: dict, layout: str) -> set:
+    """The tensor ids of the graph weights stored KRSC: every reader of the weight is a
+    convolution that takes it (`conv_weight_krsc`); a weight any other op reads keeps the order
+    that op expects."""
+    tensors, ops = dag.get("tensors") or {}, dag.get("ops") or {}
+    taken, refused = set(), set()
+    for uid in dag.get("execution_order") or []:
+        o = ops.get(uid) or {}
+        ins = o.get("input_tensor_ids") or []
+        if o.get("op_type") == "aten::convolution" and len(ins) > 1:
+            at = o.get("attributes") or {}
+            meta = tensors.get(ins[1]) or {}
+            ok = conv_weight_krsc(meta.get("shape") or [], int(at.get("groups", 1)),
+                                  bool(at.get("transposed", False)), layout)
+            (taken if ok else refused).add(ins[1])
+            rest = ins[:1] + ins[2:]
+        else:
+            rest = ins
+        refused.update(rest)
+    return {t for t in taken - refused if t.startswith(("param::", "buffer::"))}
+
+
+def is_krsc_weight(w) -> bool:
+    """A logical (K, C, R, S) weight whose memory is (K, R, S, C), contiguous."""
+    if len(w.shape) != 4:
+        return False
+    K, C, R, S = w.shape
+    return tuple(w.stride()) == (R * S * C, 1, S * C, C)
+
+
+def unit_flash(D: int, q: NBXDtype, k: NBXDtype, v: NBXDtype,
+               unit: Optional[dict] = None) -> Optional[Tuple[str, dict]]:
+    """(representation, tile) of this attention's flash on the profile's matrix unit, or None: the DtypeEngine's
+    `matrix_unit_flash_representation` ("native": the three operands are the unit's operand dtype; "split": they are
+    among `fp32_split.operands`, carried as hi + lo) and the first row of `matrix_unit.flash` — or of
+    `matrix_unit.fp32_split.flash` for the split — holding head dim D. The wrapper launches with this function.
+    `unit`: a matrix unit already read (`_configs.matrix_unit(vendor_profile)`: Prism, for the device it plans); the
+    profile in force's when None."""
+    from nbxtensor.ops._configs import matrix_unit, matrix_unit_flash_tile
+    from nbxtensor.dtype.operands import matrix_unit_flash_representation   # the library's operand rules
+    mu = matrix_unit() if unit is None else unit
+    rep = matrix_unit_flash_representation(q.name, k.name, v.name, mu)
+    if rep is None:
+        return None
+    tile = matrix_unit_flash_tile(D, mu["fp32_split"] if rep == "split" else mu)
+    return (rep, tile) if tile is not None else None
+
+
+def unit_flash_takes(D: int, q: NBXDtype, k: NBXDtype, v: NBXDtype, unit: Optional[dict] = None) -> bool:
+    """The profile's matrix unit runs this attention's flash (`unit_flash`) — the wrappers' route, the derived
+    census and Prism's priced scratch ask this same function."""
+    return unit_flash(D, q, k, v, unit) is not None
+
+
+def unit_flash_launches(Tq: int, Tk: int, D: int, q: NBXDtype, k: NBXDtype, v: NBXDtype,
+                        bias: Optional[NBXDtype] = None) -> List[Launch]:
+    """The flash on the profile's matrix unit (`unit_flash`, ops/flash_attention_m8n8k4.py): its autotune key — the
+    query rows' bucket, the keys' bucket, the head dim, the split representation — then, as Triton appends them, the
+    dtypes of its six tensors (Q, K, V, the bias — the mask's, or Q's for the zero / causal bias the wrapper builds —,
+    the output like Q, the fp32 log-sum-exp). Its tile (query rows, keys per step, warps) is certified per key
+    (workstream 1: the profile's hand table of tiles per head dim retired). No key where the unit does not take it."""
+    got = unit_flash(D, q, k, v)
+    if got is None:
+        return []
+    return [(UNIT_FLASH, (bucket_of("M", Tq), bucket_of("N", Tk), int(D), got[0] == "split",
+                          tag(q), tag(k), tag(v), tag(bias if bias is not None else q), tag(q), tag(F32)))]
+
+
+def sdpa_route(batch: int, nheads: int, Tq: int, Tk: int, D: int, Dv: int, budget_bytes: int,
+               min_chunk_rows: int, max_chunks: int, force_math: bool = False,
+               unit_flash: bool = False) -> Tuple[str, int]:
+    """The attention route: ("math", 0), ("chunked", rows) or ("flash", 0). Math when forced or
+    when the value head dim differs; else by the fp32 scores' size against the executing device's
+    budget — a non-power-of-two head dim routes on `sdpa.non_pow2_head_scores_bytes` (config/
+    tiling.yml) when the arch declares none and has not measured the flash at every head dim (`flash_any_head_dim`);
+    over the bound, flash when the profile's matrix unit takes it (`unit_flash_takes`: the
+    deterministic m8n8k4 kernel), else chunked when the rows fit the arch's ceiling; otherwise flash."""
+    if force_math or Dv != D:
+        return ("math", 0)
+    scores = _ts.sdpa_scores_bytes(batch, nheads, Tq, Tk)
+    bound = _ts.sdpa_scores_bound(budget_bytes, _pow2(D) or flash_any_head_dim())
+    if scores <= bound:
+        return ("math", 0)
+    if unit_flash:
+        return ("flash", 0)
+    if bound:
+        rows = sdpa_chunk_rows(bound, batch, nheads, Tq, Tk, min_chunk_rows, max_chunks)
+        if rows:
+            return ("chunked", rows)
+    return ("flash", 0)
+
+
+def math_attention_launches(B: int, H: int, Hk: int, Tq: int, Tk: int, D: int, Dv: int,
+                            q: NBXDtype, k: NBXDtype, v: NBXDtype, native_bf16: bool,
+                            force_accum: bool = False, contract_half: bool = False) -> List[Launch]:
+    """`_math_attention`: scores = bmm(q as [B*Hk, groups*Tq, D], k^T) — q in its (possibly
+    rounded) dtype, k in its own; then p = softmax(scores) cast to v's dtype and bmm(p, v)."""
+    M = (H // Hk) * Tq
+    # the scores are the first product's output (`bmm_out`: on the matrix unit the operands as they are in memory,
+    # elsewhere widened — fp32 for any half input without the fp16 contract, fp16 on the unit under it): the route's
+    # softmax runs on B * H * Tq rows of Tk columns at that dtype (`wrappers.softmax`, workstream 1)
+    scores = bmm_out(M, q, k, native_bf16, force_accum, contract_half)
+    return (bmm_launches(M, D, Tk, q, k, native_bf16, force_accum, contract_half)
+            + softmax_launches(B * H * Tq, Tk, False, scores, scores)
+            + bmm_launches(M, Tk, Dv, v, v, native_bf16, force_accum, contract_half))
+
+
+def chunked_math_attention_launches(B: int, H: int, Hk: int, Tq: int, Tk: int, D: int, Dv: int,
+                                    q: NBXDtype, k: NBXDtype, v: NBXDtype, native_bf16: bool,
+                                    rows: int, force_accum: bool = False, contract_half: bool = False) -> List[Launch]:
+    """`_math_attention_chunked`: the math route over query-row chunks of `rows`, the last one the
+    remainder — each chunk forms its own pair of keys."""
+    out: List[Launch] = []
+    for c in {min(rows, Tq)} | ({Tq % rows} if Tq % rows else set()):
+        out += math_attention_launches(B, H, Hk, c, Tk, D, Dv, q, k, v, native_bf16, force_accum, contract_half)
+    return out
+
+
+# ---------------------------------------------------------------------------------------------
+# Convolutions: conv2d (and conv1d as H=1, conv3d as kt conv2d launches), depthwise, bands.
+# ---------------------------------------------------------------------------------------------
+
+_NARROWEST_FIRST = (F16, BF16, F32)
+
+
+def conv_width_key(in_h: int, out_h: int, in_w: int, out_w: int, kw: int, stride_w: int, pad_w: int,
+                   dil_w: int) -> Tuple[int, int]:
+    """The width a convolution's key carries: a ONE-ROW convolution (a 1-D conv over a sequence)
+    keys its width on the profile's W ladder, the output width derived from the input's top by the
+    convolution's arithmetic; a 2-D convolution keeps its exact extents."""
+    if int(in_h) == 1 and int(out_h) == 1:
+        top = int(bucket_of("W", int(in_w)))
+        return top, (top + 2 * int(pad_w) - int(dil_w) * (int(kw) - 1) - 1) // int(stride_w) + 1
+    return int(in_w), int(out_w)
+
+
+def conv_dtypes(x: NBXDtype, w: NBXDtype) -> Tuple[NBXDtype, NBXDtype]:
+    """A convolution's operands, narrowed to the narrowest of the pair when they differ."""
+    if x != w:
+        n = next(d for d in _NARROWEST_FIRST if d in (x, w))
+        return n, n
+    return x, w
+
+
+def conv_out_hw(in_h, in_w, kh, kw, sh, sw, ph, pw, dh, dw) -> Tuple[int, int]:
+    return ((in_h + 2 * ph - dh * (kh - 1) - 1) // sh + 1,
+            (in_w + 2 * pw - dw * (kw - 1) - 1) // sw + 1)
+
+
+def conv2d_route(N: int, in_c: int, out_c: int, out_h: int, out_w: int, dh: int, dw: int, groups: int,
+                 x: NBXDtype, band_bytes: int, depthwise_enabled: bool = True) -> str:
+    """"depthwise" (groups == in_c == out_c at dilation 1), "band" (the output, sized at the
+    input's dtype, over `band_bytes`) or "plain" — the order the wrapper tests them in."""
+    from nbxtensor.dtype.nbx_dtype import dtype_size
+    if depthwise_enabled and groups == in_c and groups == out_c and dh == 1 and dw == 1:
+        return "depthwise"
+    if N * out_c * out_h * out_w * dtype_size(x) > band_bytes:
+        return "band"
+    return "plain"
+
+
+class ConvRowOverBand(ValueError):
+    """A convolution whose single output row is over the band budget: band streaming cannot split
+    it, and the per-band recursion would call itself on the same row without end."""
+
+
+def conv2d_band_rows(N: int, out_c: int, out_h: int, out_w: int, out_bytes: int, band_bytes: int) -> int:
+    """The output rows of one band of `conv2d_wrapper`'s band streaming (`_conv2d_band_streamed`):
+    bands of about `conv2d.band_budget_divisor` of the budget, evenly cut, sized by the
+    TilingEngine (`tiling_sizes.conv2d_band_rows`). A band that cannot be smaller than the output —
+    one row already over the budget — is refused by name: the recursion it would start calls the
+    same row forever (a 1-row codec conv, found by the derived census 2026-09-29)."""
+    band_oh, row_bytes = _ts.conv2d_band_rows(N, out_c, out_h, out_w, out_bytes, band_bytes)
+    if band_oh >= out_h:
+        raise ConvRowOverBand(
+            f"conv2d output [{N}, {out_c}, {out_h}, {out_w}]: one row is {row_bytes} bytes, over "
+            f"the band budget {band_bytes}; band streaming cannot split it")
+    return band_oh
+
+
+def conv2d_launches(N: int, in_c: int, in_h: int, in_w: int, out_c: int, kh: int, kw: int,
+                    sh: int, sw: int, ph: int, pw: int, dh: int, dw: int, groups: int,
+                    x: NBXDtype, w: NBXDtype, compute: Optional[NBXDtype], band_bytes: int,
+                    depthwise_enabled: bool = True, unit: Optional[dict] = None) -> List[Launch]:
+    """`conv2d_wrapper` on a 4-D weight: narrowing; the depthwise stencil when groups == in_c ==
+    out_c at dilation 1; band streaming (per-band recursion at the original padding) when the
+    output, sized at the INPUT's dtype, exceeds `band_bytes`; otherwise one `conv2d_forward_kernel`
+    launch. The output dtype is the run's compute dtype when set (`_NBX_COMPUTE_DTYPE`), else the
+    input's. The `fp16` key flag names an fp16 input (fixed 2026-09-29, aaad1c48; the entries were
+    re-keyed in place)."""
+    from nbxtensor.dtype.nbx_dtype import dtype_size
+    out_h, out_w = conv_out_hw(in_h, in_w, kh, kw, sh, sw, ph, pw, dh, dw)
+    if out_h <= 0 or out_w <= 0:
+        return []
+    x, w = conv_dtypes(x, w)
+    out = compute if compute is not None else x
+    route = conv2d_route(N, in_c, out_c, out_h, out_w, dh, dw, groups, x, band_bytes, depthwise_enabled)
+    if route == "depthwise":
+        iw_k, ow_k = conv_width_key(in_h, out_h, in_w, out_w, kw, sw, pw, 1)
+        return [(DEPTHWISE, (in_c, in_h, iw_k, out_h, ow_k, kh, kw, sh, sw, ph, pw, x == F16,
+                             tag(x), tag(w), tag(out)))]
+    xb = dtype_size(x)
+    if route == "band":
+        band_oh = conv2d_band_rows(N, out_c, out_h, out_w, xb, band_bytes)
+        out_launches: List[Launch] = []
+        for oh0 in range(0, out_h, band_oh):
+            oh1 = min(oh0 + band_oh, out_h)
+            ih0 = max(0, oh0 * sh - ph)
+            ih1 = min(in_h, (oh1 - 1) * sh + dh * (kh - 1) + 1 - ph)
+            if ih1 <= ih0:
+                continue
+            for l in conv2d_launches(N, in_c, ih1 - ih0, in_w, out_c, kh, kw, sh, sw, ph, pw, dh, dw,
+                                     groups, x, w, compute, band_bytes, depthwise_enabled, unit):
+                if l not in out_launches:
+                    out_launches.append(l)
+        return out_launches
+    if matrix_unit_conv_operands(x, w, N * out_h * out_w, out_c // max(1, groups), kh, kw, unit) is not None:  # on the matrix unit: no key
+        return []
+    iw_k, ow_k = conv_width_key(in_h, out_h, in_w, out_w, kw, sw, pw, dw)
+    return [(CONV2D, (N, in_c, in_h, iw_k, out_c, out_h, ow_k, kh, kw, sh, sw, ph, pw, dh, dw, groups,
+                      x == F16, tag(x), tag(w), tag(out)))]
+
+
+def tiled_conv2d_bands(IH: int, out_h: int, kh: int, sh: int, dh: int, pad_h: int, tile_factor: int):
+    """The bands of Prism's op-level tiled conv — `tiling_sizes.tiled_conv2d_bands`, the one cut the run, the census
+    and Prism's price of the band read."""
+    return _ts.tiled_conv2d_bands(IH, out_h, kh, sh, dh, pad_h, tile_factor)
+
+
+def tiled_conv2d_launches(N: int, in_c: int, IH: int, IW: int, out_c: int, kh: int, kw: int,
+                          sh: int, sw: int, ph: int, pw: int, dh: int, dw: int, groups: int,
+                          x: NBXDtype, w: NBXDtype, compute: Optional[NBXDtype], band_bytes: int,
+                          tile_factor: int) -> List[Launch]:
+    """The launches of an op-level tiled conv: per band of `tiled_conv2d_bands`, the band's rows
+    (padded by the image-edge rows, and by pad_w on both sides whenever any padding applies) run
+    through `conv2d_wrapper` at padding 0."""
+    out_h = (IH + 2 * ph - dh * (kh - 1) - 1) // sh + 1
+    out: List[Launch] = []
+    for _o0, _o1, start, end, pt, pb, _ht in tiled_conv2d_bands(IH, out_h, kh, sh, dh, ph, tile_factor):
+        padded = pt > 0 or pb > 0 or pw > 0
+        bh = end - start + pt + pb
+        bw = IW + 2 * pw if padded else IW
+        for l in conv2d_launches(N, in_c, bh, bw, out_c, kh, kw, sh, sw, 0, 0, dh, dw, groups,
+                                 x, w, compute, band_bytes):
+            if l not in out:
+                out.append(l)
+    return out
+
+
+def conv_launches(x_shape, w_shape, stride, padding, dilation, transposed: bool, groups: int,
+                  x: NBXDtype, w: NBXDtype, compute: Optional[NBXDtype], band_bytes: int,
+                  depthwise_enabled: bool = True, native_bf16: Optional[bool] = None, contract_half: bool = False,
+                  output_padding=0) -> List[Launch]:
+    """`conv2d_wrapper` by the weight's rank: a transposed convolution forms no autotuned key; a
+    3-D weight is conv1d as conv2d at H = 1; a 5-D weight is conv3d as kt conv2d launches over
+    [B*T_out, Cin, H, W] (identical keys). REPRODUCED as the census forms it: the chunked conv3d
+    variant is gated by the driver's free bytes, which a census answers -1 — a live run can chunk
+    and form keys this derivation does not (named)."""
+    def _n(v, n):
+        v = list(v) if isinstance(v, (list, tuple)) else [v]
+        return v + [v[-1]] * (n - len(v)) if len(v) < n else v[:n]
+    if transposed:
+        # a 1-D transposed conv of groups 1, dilation 1, no output padding runs a GEMM x^T [L, C_in] @ W [C_in, C_out * K]
+        # per batch sample (`conv_wrappers._conv_transpose1d_gemm`) and an overlap-add with no key; any other keeps the
+        # scatter kernel, which has no key
+        if len(w_shape) == 3 and groups == 1 and _n(dilation, 1)[0] == 1 and _n(output_padding, 1)[0] == 0:
+            if native_bf16 is None:
+                raise ValueError("ZERO FALLBACK: conv_launches of a 1-D transposed conv needs native_bf16 (its GEMM's key)")
+            narrowest = next(d for d in (F16, NBXDtype.bfloat16, F32) if d in (x, w)) if x != w else x
+            return mm_launches(int(x_shape[2]), int(w_shape[0]), int(w_shape[1]) * int(w_shape[2]), narrowest, narrowest,
+                               native_bf16, contract_half=contract_half)
+        return []
+    if len(w_shape) == 3:
+        N, C, L = x_shape
+        Co, Cg, K = w_shape
+        s, p, d = _n(stride, 1)[0], _n(padding, 1)[0], _n(dilation, 1)[0]
+        return conv2d_launches(N, C, 1, L, Co, 1, K, 1, s, 0, p, 1, d, groups, x, w, compute,
+                               band_bytes, depthwise_enabled)
+    if len(w_shape) == 5:
+        B, Cin, T, H, W = x_shape
+        Cout, Cg, kt, kh, kw = w_shape
+        st, sh, sw = _n(stride, 3)
+        pt, ph, pw = _n(padding, 3)
+        dt, dh, dw = _n(dilation, 3)
+        T_out = (T + 2 * pt - dt * (kt - 1) - 1) // st + 1
+        if T_out <= 0:
+            return []
+        # temporal stride 1: every tap read and written in place, ONE conv2d per batch sample over its frames
+        # [T_out, Cin, H, W] (`wrappers.conv3d_wrapper`, the taps-in-place path) — the key's batch is T_out; other
+        # strides: the frames of every sample stacked, [B*T_out, Cin, H, W]. Keyed B*T_out for both, a guidance-batch
+        # conv3d (SANA-Video's patch embed, B=2) formed a key the census never derived (W1+Apple gate, 2026-10-22).
+        return conv2d_launches(T_out if st == 1 else B * T_out, Cin, H, W, Cout, kh, kw, sh, sw, ph, pw, dh, dw,
+                               groups, x, w, compute, band_bytes, depthwise_enabled)
+    N, C, H, W = x_shape
+    Co, Cg, kh, kw = w_shape
+    sh, sw = _n(stride, 2)
+    ph, pw = _n(padding, 2)
+    dh, dw = _n(dilation, 2)
+    return conv2d_launches(N, C, H, W, Co, kh, kw, sh, sw, ph, pw, dh, dw, groups, x, w, compute,
+                           band_bytes, depthwise_enabled)
+
+
+#: The longest row the single-program softmax holds (a row's padded length in registers); a longer one runs the looped
+#: kernel (`softmax_forward_looped_kernel`, one tile at a time, keyed as the single-program one). Moved here from the
+#: wrapper so the wrapper and the census read ONE value.
+SOFTMAX_FUSED_MAX_FEAT = 16384
+
+
+def softmax_feat_bucket(feat: int) -> int:
+    """The row length's class in a softmax key: the padded row the program holds (`BLOCK_SIZE_FEAT`, the next power of
+    two) — every config's tile, its register budget and its candidates' pruning are those of the padded row, so one
+    certified entry serves every length that pads to it (keyed by the exact length, a decoder's KV positions made one
+    key per position: 11 504 keys on 16 GB, 2026-10-20)."""
+    return 1 << max(0, int(feat) - 1).bit_length()
+
+
+def softmax_launches(rows: int, feat: int, log: bool, x: NBXDtype, out: NBXDtype) -> List[Launch]:
+    """`softmax` / `log_softmax` over a row of `feat` (rows = every other element): the single-program kernel's key —
+    the rows' bucket (`bucket_of("M")`), the padded row length (`softmax_feat_bucket`), the log flag, the input's and the output's dtype;
+    a row longer than the single program holds (`SOFTMAX_FUSED_MAX_FEAT`) the looped kernel's, keyed alike (workstream
+    1: rows per program and warps certified per key; chantier 1: the looped kernel's tile and warps)."""
+    if rows < 1:
+        return []
+    kernel = SOFTMAX_LOOPED if feat > SOFTMAX_FUSED_MAX_FEAT else SOFTMAX
+    return [(kernel, (bucket_of("M", rows), softmax_feat_bucket(feat), bool(log), tag(x), tag(out)))]
+
+
+def layer_norm_launches(rows: int, feat: int, has_weight: bool, has_bias: bool, x: NBXDtype,
+                        w: Optional[NBXDtype], b: Optional[NBXDtype], out: NBXDtype) -> List[Launch]:
+    """`wrappers.native_layer_norm` over rows of `feat` (the normalized extent; every other element a row): the key —
+    the rows' bucket (`bucket_of("M")`), the padded row (`softmax_feat_bucket`: the program holds the next power of
+    two, one certified entry per padding), the weight and bias flags (a bias without a weight is not applied, as the
+    wrapper reads it), the statistics always saved — then the dtype tags of its tensor arguments in order: input,
+    weight (the input standing in when there is none), bias (likewise), mean and inverse deviation (fp32), output."""
+    if rows < 1 or feat < 1:
+        return []
+    bias = bool(has_bias and has_weight)
+    wt = w if has_weight else x
+    bt = b if bias else x
+    return [(LAYER_NORM, (bucket_of("M", rows), softmax_feat_bucket(feat), bool(has_weight), bias, True,
+                          tag(x), tag(wt), tag(bt), tag(F32), tag(F32), tag(out)))]
+
+
+def rms_norm_launches(rows: int, feat: int, has_weight: bool, x: NBXDtype, w: Optional[NBXDtype],
+                      out: NBXDtype) -> List[Launch]:
+    """`wrappers.rms_norm` over rows of `feat` (the last dim): the key — the rows' bucket, the padded row, the weight
+    flag, the protected fp16 store (`nbx_tensor._saturating_cast` from fp32, the kernel's accumulator, to the output's
+    dtype) — then the dtype tags of input, weight (the input standing in when there is none), output."""
+    if rows < 1 or feat < 1:
+        return []
+    from nbxtensor.tensor import _saturating_cast
+    return [(RMS_NORM, (bucket_of("M", rows), softmax_feat_bucket(feat), bool(has_weight),
+                        bool(_saturating_cast(F32, out)), tag(x), tag(w if has_weight else x), tag(out)))]
+
+
+def elementwise_classes(largest: int) -> List[Launch]:
+    """Every element-wise class (`ops/elementwise_launch.ew_bucket`: the profile's `elements` ladder, else its
+    `default`) a component can launch: each class whose floor is at most `largest`, the most elements any tensor of
+    the component holds — an element-wise launch writes or reads at most one tensor's elements (an op's output, a
+    copy, a cast, a weight converted at load). Keyed as the shared copy is (`ew_launch_kernel`: the class's top and
+    its two fp32 operands)."""
+    from nbxtensor.ops.autotune_bucket import ladder_for
+    largest = max(1, int(largest))
+    out, prev = [], 0
+    for bound, step in ladder_for("elements"):
+        hi = largest if bound is None else min(bound, largest)
+        top = -(-(prev + 1) // step) * step
+        while prev < hi and top - step < hi:
+            out.append((ELEMENTWISE, (top, "fp32", "fp32")))
+            top += step
+            if bound is not None and top > bound:
+                break
+        if bound is None or largest <= bound:
+            break
+        prev = bound
+    return out
+
+
+def group_norm_launches(N: int, C: int, HxW: int, groups: int, has_weight: bool, has_bias: bool, x: NBXDtype,
+                        w: Optional[NBXDtype], b: Optional[NBXDtype]) -> List[Launch]:
+    """`wrappers.group_norm_wrapper` (one program per batch x group, the group's channel x spatial extent walked in
+    chunks): the key — the programs' bucket (`bucket_of("M")`), the group's padded extent (`softmax_feat_bucket`),
+    the weight and bias flags — then the dtype tags of its tensor arguments in order: input, output (the input's),
+    weight and bias (the input standing in for an absent one), mean and inverse deviation (fp32)."""
+    if N < 1 or C < 1 or HxW < 1 or groups < 1:
+        return []
+    hidden = (C // groups) * HxW
+    return [(GROUP_NORM, (bucket_of("M", N * groups), softmax_feat_bucket(hidden), bool(has_weight), bool(has_bias),
+                          tag(x), tag(x), tag(w if has_weight else x), tag(b if has_bias else x),
+                          tag(F32), tag(F32)))]
+
+
+def row_reduce_launches(kind: str, rows: int, feat: int, x: NBXDtype) -> List[Launch]:
+    """`wrappers.mean_wrapper` / `sum_wrapper` / `amax_wrapper` over rows of `feat` (the reduced elements: one dim, a
+    list flattened by `_reduce_over_dims`, or all of them; every other element a row): the key — the rows' bucket
+    (`bucket_of("M")`), the padded row (`softmax_feat_bucket`) — then the input's dtype tag (after the `dtype` keyword's
+    cast) and the output's: fp32 for mean and sum (cast after), the input's for amax."""
+    if rows < 1 or feat < 1:
+        return []
+    out = x if kind == "amax" else F32
+    return [(ROW_REDUCE[kind], (bucket_of("M", rows), softmax_feat_bucket(feat), tag(x), tag(out)))]
+
+
+def fused_moe_launches(tokens: int, top_k: int, experts: int, hidden: int, inter: int, w: NBXDtype) -> List[Launch]:
+    """A fused MoE layer's three grouped GEMMs (`moe_wrappers.invoke_fused_moe`, `triton/moe._fused_moe_pass` and its
+    multi-card and expert-store forms): gate and up [hidden -> inter], down [inter -> hidden] with the routing weight
+    and no top-k division. Every float operand is in the experts' dtype `w` (the layer casts the hidden states and the
+    scores to it), the routing tables int64. The key's length class is the valid (token, expert) pairs' bucket,
+    `tokens * top_k` — the same for every launch of the layer, whatever the routing (the padded extent is not keyed) —
+    and `experts` the pointer table's size (the layer's experts)."""
+    if tokens < 1:
+        return []
+    nv = bucket_of("M", int(tokens) * int(top_k))
+    t, i64 = tag(w), tag(NBXDtype.int64)
+    dts = (t, i64, t, t, i64, i64, i64)        # a, expert_ptrs, c, topk_weights, sorted ids, expert ids, padded count
+    return [(FUSED_MOE, (int(inter), int(hidden), nv, int(experts), False, True) + dts),
+            (FUSED_MOE, (int(hidden), int(inter), nv, int(experts), True, False) + dts)]
+
+
+def moe_align_classes(experts: int, pairs: int) -> Tuple[int, int]:
+    """The MoE alignment's padded expert count (a power of two, at least 16) and its key's length class: the layer's
+    (token, expert) pairs in the profile's `elements` ladder — the alignment is integer bookkeeping over a few
+    thousand ids, launch-dominated as an element-wise launch is (the ladder's own measurement: at or under 65 536
+    elements the launch dominates whatever the tile), so its classes are those few, certified at their floor
+    (`autotune_bucket.class_floor`); the GEMM's row ladder made ~300 classes of it a memory class."""
+    from triton import next_power_of_2
+    return max(next_power_of_2(int(experts)), 16), bucket_of("elements", int(pairs))
+
+
+def moe_align_launches(tokens: int, top_k: int, experts: int, w: NBXDtype) -> List[Launch]:
+    """The three alignment stages a fused MoE layer launches (`triton/moe.moe_align_block_size`), the routing ids
+    int64: at the FMA kernel's block (`moe_wrappers._MOE_BM`) and, on a matrix unit that takes the experts' dtype,
+    at every unit tile's rows (`fused_moe_unit_tiles` — the layer aligns again at the tile its certificate chose,
+    one of these)."""
+    if tokens < 1:
+        return []
+    from nbxtensor.ops.moe_wrappers import _MOE_BM
+    from nbxtensor.ops._configs import matrix_unit
+    be, nb = moe_align_classes(experts, int(tokens) * int(top_k))
+    e, i64 = int(experts), tag(NBXDtype.int64)
+    blocks = [_MOE_BM]
+    if w == F16:
+        blocks += sorted({int(r["block_m"]) for r in ((matrix_unit() or {}).get("mm") or [])} - {_MOE_BM})
+    out: List[Launch] = []
+    for bs in blocks:
+        out += [(MOE_ALIGN[0], (bs, be, e, nb, i64, i64, i64, i64)),
+                (MOE_ALIGN[1], (bs, be, e, nb, i64, i64, i64, i64))]
+    return out + [(MOE_ALIGN[2], (e, nb, i64, i64, i64))]
+
+
+def moe_decode_vec_half_launches(top_k: int, hidden: int, inter: int, w: NBXDtype) -> List[Launch]:
+    """A one-token MoE layer on half-precision experts held on one card (`triton/moe._moe_decode_vec_half_pass`): the
+    gateup with SwiGLU [hidden -> inter], the down projection's router-weighted partials [inter -> hidden], their
+    fixed-order sum. The routed expert ids int64, the tables int64, the partials fp32; the activation and the router
+    weights in the experts' dtype `w`."""
+    if w not in (NBXDtype.float16, NBXDtype.bfloat16):
+        return []
+    t, i64, f32 = tag(w), tag(NBXDtype.int64), tag(F32)
+    bf = w == NBXDtype.bfloat16
+    return [(MOE_GATEUP_HALF, (int(hidden), int(inter), int(top_k), bf, t, i64, i64, i64, f32)),
+            (MOE_DOWN_HALF, (int(inter), int(hidden), int(top_k), bf, f32, i64, t, i64, f32)),
+            (MOE_PART_REDUCE, (int(hidden), int(top_k), f32, f32))]
+
+
+def decode_reduce_launch(splits: int, Dv: int, out: NBXDtype) -> Launch:
+    """`attention_wrappers.decode_reduce`'s launch, the fixed-order combine after both split decodes: keyed by the
+    splits, one row per program (GROUPS 1) and the value head dim; the partials fp32, the output in `out`."""
+    f32 = tag(F32)
+    return (DECODE_REDUCE, (int(splits), 1, int(Dv), f32, f32, f32, tag(out)))
+
+
+def decode_vec_launches(H: int, Hk: int, Tk: int, D: int, Dv: int, q: NBXDtype, k: NBXDtype, v: NBXDtype,
+                        bias: Optional[NBXDtype], q_round: Optional[NBXDtype] = None) -> List[Launch]:
+    """`attention_wrappers._decode_attn_vec`'s split launch: one query row per program, keyed by the keys' bucket, the
+    head dims, the GQA group and the bias; Q as it arrives (the kernel rounds it to `q_round`, the cache's dtype, on
+    load), the partials fp32; then the shared reduce over its splits (`decode_vec_splits`), written in `q_round`, else
+    in Q's dtype."""
+    from nbxtensor.ops.decode_attn_vec import decode_vec_splits
+    f32 = tag(F32)
+    return [(DECODE_VEC, (bucket_of("N", Tk), int(D), int(Dv), int(H) // max(1, int(Hk)), bias is not None,
+                          tag(q), tag(k), tag(v), tag(bias if bias is not None else q), f32, f32, f32)),
+            decode_reduce_launch(decode_vec_splits(Tk), Dv, q_round if q_round is not None else q)]
+
+
+def paged_decode_launches(H: int, Hk: int, Tk: int, D: int, Dv: int, q: NBXDtype, kv: NBXDtype,
+                          page_tokens: int) -> List[Launch]:
+    """`attention_wrappers.paged_decode_attention`'s split launch (a decode step through the paged KV cache, the Triton
+    engine's every cached decode): keyed by the splits the longest sequence makes (`paged_decode_splits`, the replayed
+    decode's bucket too), the page, the head dims and the GQA group;
+    Q as it arrives (rounded to the cache's dtype on load), the pools in the cache's, the tables int32, the partials
+    fp32; then the shared reduce over those splits, written in the cache's dtype (the cache's decode rounds Q to it,
+    `kv_cache.intercept`)."""
+    f32, i32, c = tag(F32), tag(NBXDtype.int32), tag(kv)
+    from nbxtensor.ops.paged_decode_attn import paged_decode_splits
+    return [(PAGED_DECODE, (paged_decode_splits(Tk), int(page_tokens), int(D), int(Dv), int(H) // max(1, int(Hk)),
+                            tag(q), c, c, i32, i32, f32, f32, f32)),
+            decode_reduce_launch(paged_decode_splits(Tk), Dv, kv)]
+