nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/matmul.py

@@ -0,0 +1,377 @@
+"""Matrix multiplication — pure @triton.jit kernel with autotune.
+
+Ported from Triton official tutorial (BSD license)
+https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+adapted for NeuroBrix:
+  - NBXTensor inputs (handled by wrappers, kernel sees raw ptrs/strides)
+  - Volta-aware static dtype path (PROMOTE_B + IEEE_PRECISION constexpr)
+  - @triton.autotune across 18 configs (Phase 1.5, 2026-05): the only
+    proven path to ≥70% cuBLAS HMMA on Sana DiT shapes — see CLAUDE.md
+    "Autotune policy" section for the doctrinal exception that allows
+# PROMOTE_A (and PROMOTE_BIAS) are NOT part of the key: the directory's certified
+# settings are indexed by the shape and the pre-existing flags, and a promotion of
+# an operand on load changes neither the tiles' legality nor the accumulation
+# order — the same setting applies; keying on it made every matmul of the tree a
+# miss (5,628 entries unserved, a runtime sweep per key, 2026-09-07 lever 2).
+    @triton.autotune on mm/bmm/addmm/conv2d.
+  - tl.dot 3-arg HMMA-FMA fused form
+  - tl.assume integer-analyzer hints
+  - in-kernel cast accum → output dtype
+Handles mm (2D) + addmm (with bias). bmm handled in wrappers via batch loop.
+"""
+
+import triton
+
+from nbxtensor.ops._configs import arch_smem_budget, configs_within_smem_budget, nbx_autotune
+import triton.language as tl
+
+from nbxtensor.ops._autotune_policy import maybe_pin_single, is_matmul_pinned
+from nbxtensor.ops._common import sigmoid, tanh_fn
+
+# Vertical-fusion epilogue codes (fusion_vertical pass, EXACT gate).
+# 0 = none, 1 = silu, 2 = gelu exact, 3 = gelu tanh. The epilogue is a
+# tl.constexpr NOT in the autotune key: both fused and unfused variants
+# share ONE tuning-cache entry per (M, N, K, IEEE, PROMOTE_B), so the
+# selected config — and therefore the tl.dot accumulation order — is
+# identical by construction (the two-leg byte-exactness contract: the
+# R16 research showed rounding emulation alone is insufficient if the
+# fused and unfused kernels can tune to different BLOCK_K).
+#
+# Byte-exactness emulation: the unfused pair is
+#   matmul_kernel: fp32 acc -> cast to C dtype -> store        (round 1)
+#   silu/gelu kernel: load C dtype -> .to(fp32) -> f(x) -> store (round 2)
+# The fused epilogue reproduces BOTH rounding points in-register:
+#   acc -> .to(C dtype) -> .to(fp32) -> f(x) -> store
+# using the SAME formula text and the SAME _common helpers as the
+# standalone kernels, so the arithmetic lowers identically.
+EPILOGUE_NONE: int = 0
+EPILOGUE_SILU: int = 1
+EPILOGUE_GELU_EXACT: int = 2
+EPILOGUE_GELU_TANH: int = 3
+
+
+# Per-architecture autotune configs (tutorial pattern
+# `get_cuda_autotune_config()` adapted to NeuroBrix: gate by detected
+# compute capability so each hardware explores ONLY its viable subspace).
+# Volta sm_70 has 96 KB SMEM/SM (vs 192 KB sm_80 / 228 KB sm_90); large
+# blocks (BM≥128 BK≥64 warps=8) saturate SMEM → register spill →
+# catastrophic perf (98-145 ms measured Phase 1.5 Step 2 FlagGems
+# bench). The Volta-viable subspace is restricted to BM∈{32,64},
+# BN∈{32,64,128}, BK∈{32,64}, warps∈{2,4}, stages∈{2..5} — ~20
+# combinations giving the autotuner a denser space of fitting configs.
+# Ampere+ (sm_80+) gates BM/BN/BK larger as the hardware supports it.
+
+_MATMUL_AUTOTUNE_VOLTA = [
+    # BM=32 row
+    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32,  'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=4, num_warps=2),
+    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32,  'BLOCK_K': 64, 'GROUP_M': 8}, num_stages=4, num_warps=2),
+    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64,  'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=4, num_warps=2),
+    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64,  'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=5, num_warps=2),
+    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64,  'BLOCK_K': 64, 'GROUP_M': 8}, num_stages=3, num_warps=4),
+    triton.Config({'BLOCK_M': 32, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=3, num_warps=4),
+    # BM=64 row (the historical NeuroBrix default neighborhood)
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32,  'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=5, num_warps=2),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 32,  'BLOCK_K': 64, 'GROUP_M': 8}, num_stages=4, num_warps=2),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64,  'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64,  'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=3, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64,  'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64,  'BLOCK_K': 64, 'GROUP_M': 8}, num_stages=3, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64,  'BLOCK_K': 64, 'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=2, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=3, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64, 'GROUP_M': 8}, num_stages=3, num_warps=4),
+]
+
+# Ampere/Hopper subspace (sm_80+) — larger blocks allowed thanks to
+# SMEM doubling and tensor-core throughput. Tutorial canonical configs.
+_MATMUL_AUTOTUNE_AMPERE_PLUS = [
+    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64,  'GROUP_M': 8}, num_stages=3, num_warps=8),
+    triton.Config({'BLOCK_M': 64,  'BLOCK_N': 256, 'BLOCK_K': 32,  'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32,  'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64,  'BLOCK_K': 32,  'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 64,  'BLOCK_N': 128, 'BLOCK_K': 32,  'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64,  'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64,  'BLOCK_K': 64,  'GROUP_M': 8}, num_stages=4, num_warps=4),
+    triton.Config({'BLOCK_M': 64,  'BLOCK_N': 128, 'BLOCK_K': 64,  'GROUP_M': 8}, num_stages=4, num_warps=4),
+] + _MATMUL_AUTOTUNE_VOLTA  # Ampere+ also benefits from smaller configs
+
+
+def _detect_arch_configs():
+    """Return the autotune config space for the active architecture.
+
+    WHY THIS ASKS THE DRIVER, when R23/R24 put hardware parameters in
+    `config/vendors/<vendor>/<arch>.yml`: autotune configs are bound at
+    DECORATION time, which is module import — before `set_hardware_profile()`
+    has run at executor construction. There is no profile to read yet. This is
+    the documented bootstrap exception, not a licence to query the driver
+    elsewhere; every parameter that CAN wait for the profile reads the YAML
+    (see `_safe_num_stages`, which moved to `pipelining.max_num_stages`).
+
+    The spaces:
+
+    * sm_80+ (Ampere/Hopper) — tutorial canonical plus the Volta subset.
+    * sm_70 (Volta) — Volta-viable subset only: larger blocks saturate the
+      96 KB SMEM, spill registers and collapse throughput (measured, Phase 1.5).
+    * CDNA (gfx*) — the Volta-viable subset. AMD's Triton target reports `arch`
+      as a STRING ("gfx90a"), so the old `isinstance(cap, int)` test silently
+      dropped every AMD card into the Volta space *by accident*. It lands in the
+      same place deliberately now, and for a stated reason: CDNA exposes 64 KB
+      of LDS against Volta's 96 KB, so the small-block space is the safe one
+      until an MI-series card measures a better one. Named as a first-light
+      task, not left as a coincidence.
+    * anything else (Apple, unknown) — the Volta subset as the STARTING space.
+
+    Whichever space is chosen, it is then filtered against the executing
+    hardware's `memory.max_shared_memory_per_block`, read from its vendor YAML.
+    Choosing by architecture NAME alone was wrong in two directions and both
+    are now decided rather than inherited:
+
+    * **Apple** was sent to the Volta subset because that space "fits the
+      smallest budget". It does not: Volta declares 96 KB and an Apple GPU has
+      32 KB, so a BM=64/BN=128/BK=64 tile at 3 stages wants 72 KB and cannot
+      run there. Nine of the seventeen Volta configs survive the real budget.
+    * **CDNA** landed in the Volta space by accident (a string arch failed an
+      `isinstance(cap, int)` test), was later made deliberate, and still had no
+      budget applied — CDNA1/2 declare 64 KB, which excludes the largest Volta
+      tile at 72 KB.
+
+    Filtering never empties the space and never invents a budget: with no
+    matching profile the space is returned untouched.
+    """
+    try:
+        from nbxtensor.backend.launcher import target as _nbx_target  # engine data, no driver probe (R33)
+        arch = _nbx_target().arch
+    except Exception:
+        return _MATMUL_AUTOTUNE_VOLTA
+
+    # NVIDIA: Triton reports compute capability x 10 (sm_70 -> 70, sm_80 -> 80)
+    if isinstance(arch, int):
+        space = (_MATMUL_AUTOTUNE_AMPERE_PLUS if arch >= 80
+                 else _MATMUL_AUTOTUNE_VOLTA)
+    else:
+        # AMD CDNA/RDNA and any other string-arch target (Apple included).
+        space = _MATMUL_AUTOTUNE_VOLTA
+
+    return configs_within_smem_budget(space, arch_smem_budget())
+
+
+_MATMUL_AUTOTUNE_CONFIGS = maybe_pin_single(
+    _detect_arch_configs(), is_matmul_pinned)
+
+
+@nbx_autotune(configs=_MATMUL_AUTOTUNE_CONFIGS,
+                 key=['M_BUCKET', 'N', 'K', 'IEEE_PRECISION', 'PROMOTE_B'],
+                 cache_results=True)
+@triton.jit
+def matmul_kernel(
+    a_ptr, b_ptr, c_ptr,
+    M, N, K, M_BUCKET,
+    stride_am, stride_ak,
+    stride_bk, stride_bn,
+    stride_cm, stride_cn,
+    IEEE_PRECISION: tl.constexpr = False,
+    PROMOTE_B: tl.constexpr = False,
+    PROMOTE_A: tl.constexpr = False,
+    EPILOGUE: tl.constexpr = 0,
+    BLOCK_M: tl.constexpr = 64,
+    BLOCK_N: tl.constexpr = 64,
+    BLOCK_K: tl.constexpr = 32,
+    GROUP_M: tl.constexpr = 8,
+):
+    """C = A @ B where A is [M, K], B is [K, N], C is [M, N].
+
+    Accumulates in fp32 for numerical stability.
+    Output dtype determined by output pointer's dtype (in-kernel cast).
+
+    IEEE_PRECISION=True forces `tl.dot(input_precision="ieee")` — required
+    when fp32 inputs carry magnitudes > fp16_max on pre-Ampere GPUs,
+    because `tl.dot` otherwise lowers through fp16 HMMA tensor cores which
+    saturate the inputs to fp16 before the multiply. Set by the wrapper
+    when `not _NBX_HAS_NATIVE_BF16` and inputs were promoted to fp32.
+
+    PROMOTE_B=True casts the b tile to a's dtype after load and before
+    tl.dot. Triton's type checker rejects `tl.dot(fp32, fp16)` at compile
+    time; the cast is the cheapest way to bridge the mismatch — fused
+    with the load, register-level, no heap allocation. Set by the wrapper
+    when the activation was upcast fp16→fp32 (step 2) but the weight
+    was left fp16 in memory (to save VRAM). The bit-exact fp32 promotion
+    of a fp16 tile is free numerically (fp16 values are a subset of
+    fp32); the accumulator is fp32 so the final dot product is identical
+    to the path that widens the full weight pre-kernel.
+
+# PROMOTE_A / PROMOTE_BIAS are not part of the key — see matmul_kernel above.
+    Phase 1.5 (2026-05): @triton.autotune ENABLED. The autotune key
+    includes IEEE_PRECISION + PROMOTE_B so each (Volta-fp32 / Volta-fp16-mixed
+    / Ampere+ pure fp16) path gets its own selected config.
+    """
+    # tl.assume hints (tutorial pattern) — combined with autotune they
+    # carry through; without autotune they were inactive (rollback 396fef1).
+    tl.assume(stride_am > 0); tl.assume(stride_ak > 0)
+    tl.assume(stride_bk > 0); tl.assume(stride_bn > 0)
+    tl.assume(stride_cm > 0); tl.assume(stride_cn > 0)
+    pid = tl.program_id(0).to(tl.int64)
+    num_pid_m = tl.cdiv(M, BLOCK_M)
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    num_pid_in_group = GROUP_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
+    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
+    offs_k = tl.arange(0, BLOCK_K)
+    # Pointer offsets in int64: an int32 row offset times a stride wraps past
+    # 2^31 elements (Mochi's VAE mm, M=1 068 480 x N=2048 — CUDA 700 at every
+    # attempt since 2026-09-10; test_a_gemm_beyond_two_billion_elements).
+    a_ptrs = a_ptr + (offs_am[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :].to(tl.int64) * stride_bn)
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
+        if PROMOTE_A:
+            # The activation tile widened in registers — the same numbers the wrapper
+            # produced by materialising `a.to(float32)` before the call (an exact widening),
+            # without that copy per matmul (TinyLlama decode: 143 a token, 2026-09-07).
+            a = a.to(tl.float32)
+        if PROMOTE_B:
+            b = b.to(a.dtype)
+        # 3-arg HMMA-FMA fused form (tutorial pattern).
+        if IEEE_PRECISION:
+            accumulator = tl.dot(a, b, accumulator, input_precision="ieee")
+        else:
+            accumulator = tl.dot(a, b, accumulator)
+        a_ptrs += BLOCK_K * stride_ak
+        b_ptrs += BLOCK_K * stride_bk
+
+    # In-kernel cast accum → output dtype (tutorial pattern; faster than
+    # tl.store auto-cast since the conversion happens in registers).
+    c = accumulator.to(c_ptr.dtype.element_ty)
+    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    c_ptrs = c_ptr + stride_cm * offs_cm[:, None].to(tl.int64) + stride_cn * offs_cn[None, :].to(tl.int64)
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+    if EPILOGUE != 0:
+        # Per-stage rounding emulation: `c` already carries round 1 (the
+        # unfused mm's store cast); .to(fp32) is the standalone epilogue
+        # kernel's lossless load upcast. Formulas are copied VERBATIM
+        # from ops/silu.py / ops/gelu.py; the store's implicit downcast
+        # mirrors the standalone kernel's tl.store.
+        x_fp32 = c.to(tl.float32)
+        if EPILOGUE == 1:
+            out = x_fp32 * sigmoid(x_fp32)
+        elif EPILOGUE == 2:
+            cdf = 0.5 * (1.0 + tl.math.erf(0.707106781 * x_fp32))
+            out = cdf * x_fp32
+        else:
+            cdf = 0.5 * (1.0 + tanh_fn(0.7978845608 * x_fp32 * (1.0 + 0.044715 * x_fp32 * x_fp32)))
+            out = cdf * x_fp32
+        tl.store(c_ptrs, out, mask=c_mask)
+    else:
+        tl.store(c_ptrs, c, mask=c_mask)
+
+
+@nbx_autotune(configs=_MATMUL_AUTOTUNE_CONFIGS,
+                 key=['M_BUCKET', 'N', 'K', 'IEEE_PRECISION', 'PROMOTE_B'],
+                 cache_results=True)
+@triton.jit
+def addmm_kernel(
+    a_ptr, b_ptr, bias_ptr, c_ptr,
+    M, N, K, M_BUCKET,
+    stride_am, stride_ak,
+    stride_bk, stride_bn,
+    stride_cm, stride_cn,
+    alpha, beta,
+    IEEE_PRECISION: tl.constexpr = False,
+    PROMOTE_B: tl.constexpr = False,
+    PROMOTE_A: tl.constexpr = False,
+    PROMOTE_BIAS: tl.constexpr = False,   # the bias widened to fp32 on load (exact), not cast beforehand
+    EPILOGUE: tl.constexpr = 0,
+    BLOCK_M: tl.constexpr = 64,
+    BLOCK_N: tl.constexpr = 64,
+    BLOCK_K: tl.constexpr = 32,
+    GROUP_M: tl.constexpr = 8,
+):
+    """C = beta * bias + alpha * (A @ B) where bias is [N].
+
+    PROMOTE_B: see matmul_kernel docstring. Same in-kernel fp16→fp32
+    tile cast; enables the wrapper to keep fp16 weights fp16 in memory
+    while still running tl.dot with matched dtypes.
+
+    Phase 1.5 (2026-05): @triton.autotune ENABLED — same configs as
+    matmul_kernel. Tutorial pattern adapted with NeuroBrix's bias
+    addition + alpha/beta scaling.
+    """
+    tl.assume(stride_am > 0); tl.assume(stride_ak > 0)
+    tl.assume(stride_bk > 0); tl.assume(stride_bn > 0)
+    tl.assume(stride_cm > 0); tl.assume(stride_cn > 0)
+    pid = tl.program_id(0).to(tl.int64)
+    num_pid_m = tl.cdiv(M, BLOCK_M)
+    num_pid_n = tl.cdiv(N, BLOCK_N)
+    num_pid_in_group = GROUP_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M
+    offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N
+    offs_k = tl.arange(0, BLOCK_K)
+    # Pointer offsets in int64: an int32 row offset times a stride wraps past
+    # 2^31 elements (Mochi's VAE mm, M=1 068 480 x N=2048 — CUDA 700 at every
+    # attempt since 2026-09-10; test_a_gemm_beyond_two_billion_elements).
+    a_ptrs = a_ptr + (offs_am[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :].to(tl.int64) * stride_bn)
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
+        if PROMOTE_A:
+            # The activation tile widened in registers — the same numbers the wrapper
+            # produced by materialising `a.to(float32)` before the call (an exact widening),
+            # without that copy per matmul (TinyLlama decode: 143 a token, 2026-09-07).
+            a = a.to(tl.float32)
+        if PROMOTE_B:
+            b = b.to(a.dtype)
+        if IEEE_PRECISION:
+            accumulator = tl.dot(a, b, accumulator, input_precision="ieee")
+        else:
+            accumulator = tl.dot(a, b, accumulator)
+        a_ptrs += BLOCK_K * stride_ak
+        b_ptrs += BLOCK_K * stride_bk
+
+    # Add bias: C = alpha * matmul + beta * bias
+    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    bias_mask = offs_cn < N
+    bias = tl.load(bias_ptr + offs_cn, mask=bias_mask)
+    if PROMOTE_BIAS:
+        bias = bias.to(tl.float32)
+    accumulator = alpha * accumulator + beta * bias[None, :]
+
+    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    c_ptrs = c_ptr + stride_cm * offs_cm[:, None].to(tl.int64) + stride_cn * offs_cn[None, :].to(tl.int64)
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+    if EPILOGUE != 0:
+        # Per-stage rounding emulation (see matmul_kernel): the unfused
+        # addmm stores the fp32 accumulator with tl.store's implicit
+        # downcast (round 1); the standalone epilogue kernel loads and
+        # upcasts to fp32. Reproduce both in-register, then apply the
+        # VERBATIM standalone formulas.
+        x_fp32 = accumulator.to(c_ptr.dtype.element_ty).to(tl.float32)
+        if EPILOGUE == 1:
+            out = x_fp32 * sigmoid(x_fp32)
+        elif EPILOGUE == 2:
+            cdf = 0.5 * (1.0 + tl.math.erf(0.707106781 * x_fp32))
+            out = cdf * x_fp32
+        else:
+            cdf = 0.5 * (1.0 + tanh_fn(0.7978845608 * x_fp32 * (1.0 + 0.044715 * x_fp32 * x_fp32)))
+            out = cdf * x_fp32
+        tl.store(c_ptrs, out, mask=c_mask)
+    else:
+        tl.store(c_ptrs, accumulator, mask=c_mask)

nbxtensor/ops/matmul_m8n8k4.py

@@ -0,0 +1,235 @@
+"""GEMM on the m8n8k4 matrix unit — `mm` and `addmm` where the hardware profile declares `matrix_unit`.
+
+Same contract as `matmul_kernel` / `addmm_kernel` / `baddbmm_kernel` (ops/matmul.py) for operands that are both the unit's operand
+dtype in memory: C[z] = A[z] @ B[z] (+ beta * bias + alpha scaling; the bias broadcast by its strides; batch on
+grid axis 1, stride 0 for mm / addmm), fp32 accumulation, any strides on B (a
+pre-transposed weight is read in place), int64 offsets, M/N/K tails masked, the output cast to C's dtype, and the
+fused epilogues (silu, gelu exact, gelu tanh) with the per-stage rounding emulation of the unfused pair: the
+accumulator is rounded to C's dtype, widened back, then the standalone kernel's formula is applied.
+
+The products of two fp16 operands are exact in fp32. The unit's own accumulation rounds toward zero (Fasi et al.
+2021, PeerJ CS 7:e330), so it sums only one K tile into a zeroed accumulator; the tiles are summed on the FMA units
+with round to nearest, as the fp32 `tl.dot(input_precision="ieee")` path sums, in another order. The tile is the profile's
+(`matrix_unit.mm`), never tuned at run time: no autotune key, nothing to certify.
+
+Operands go global -> registers (coalesced) -> shared memory as [k-step x rows, 4] -> HMMA fragments
+(ops/mma_m8n8k4.py); the next K tile is loaded while the current one is multiplied.
+
+SPLIT representation (SPLIT_A / SPLIT_B, decided by the DtypeEngine: nbxtensor/dtype/operands.matrix_unit_representation, the
+profile's `matrix_unit.fp32_split`): an operand of any float dtype in memory is widened to fp32, each row of its K tile (each
+column for B) scaled by the power of two 2^s that brings that row's largest magnitude to [2^HI_EXP, 2^(HI_EXP+1)),
+and carried as hi = fp16(x 2^s) and lo = fp16((x 2^s - hi) 2^LO_SHIFT). The scale is constant along k, so it factors
+out of the dot product: a small row next to a large one keeps its own precision. Per K tile the unit computes hi*hi
+into one zeroed accumulator and hi*lo + lo*hi into another; the tile's sum t_hh + t_lo 2^-LO_SHIFT, times 2^-s of
+each split operand's row (column), is added to the fp32 accumulator on the FMA units (round to nearest). The unit's own accumulation rounds
+toward zero (Fasi et al. 2021, PeerJ CS 7:e330), so it never carries the sum across tiles; lo*lo (< 2^-22 of the
+product) is not formed (Ootomo & Yokota 2022, arXiv 2203.03341, eq. 24). Every power-of-two scaling is exact.
+"""
+
+import functools
+
+from triton.experimental import gluon
+from triton.experimental.gluon import language as gl
+
+from .mma_m8n8k4 import (acc_layout, acc_rows_cols, bcast_rows, frag_a_layout, frag_b_layout, mma_k4,
+                         stage_k_inner, stage_k_outer)
+
+
+@gluon.jit
+def _sigmoid(x):
+    return 1.0 / (1.0 + gl.exp(-x))
+
+
+@gluon.jit
+def _tanh(x):
+    return 2.0 * _sigmoid(2.0 * x) - 1.0
+
+
+@gluon.jit
+def matmul_m8n8k4_kernel(
+    A, B, Bias, C,
+    M, N, K,
+    stride_az, stride_am, stride_ak, stride_bz, stride_bk, stride_bn, stride_cz, stride_cm, stride_cn,
+    stride_xz, stride_xm, stride_xn,
+    alpha, beta,
+    HAS_BIAS: gl.constexpr, PROMOTE_BIAS: gl.constexpr, EPILOGUE: gl.constexpr, B_K_CONTIG: gl.constexpr,
+    SPLIT_A: gl.constexpr, SPLIT_B: gl.constexpr, HI_EXP: gl.constexpr, LO_SHIFT: gl.constexpr,
+    MB: gl.constexpr, NB: gl.constexpr, BK: gl.constexpr, GROUP_M: gl.constexpr,
+    L: gl.constexpr, LA: gl.constexpr, LB: gl.constexpr, GA: gl.constexpr, GB: gl.constexpr,
+):
+    BM: gl.constexpr = 16 * MB
+    BN: gl.constexpr = 16 * NB
+    J: gl.constexpr = BK // 4
+    SH: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, order=[1, 0])
+
+    pid = gl.program_id(0).to(gl.int64)
+    z = gl.program_id(1).to(gl.int64)                      # the batch (0 for mm / addmm)
+    A = A + z * stride_az
+    B = B + z * stride_bz
+    C = C + z * stride_cz
+    num_pid_m = gl.cdiv(M, BM)
+    num_pid_n = gl.cdiv(N, BN)
+    num_pid_in_group = GROUP_M * num_pid_n
+    first_pid_m = (pid // num_pid_in_group) * GROUP_M
+    group_size_m = gl.minimum(num_pid_m - first_pid_m, GROUP_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    m0 = pid_m * BM
+    n0 = pid_n * BN
+
+    sa = gl.allocate_shared_memory(gl.float16, [J * BM, 4], SH)
+    sb = gl.allocate_shared_memory(gl.float16, [J * BN, 4], SH)
+    if SPLIT_A:
+        sal = gl.allocate_shared_memory(gl.float16, [J * BM, 4], SH)
+    if SPLIT_B:
+        sbl = gl.allocate_shared_memory(gl.float16, [J * BN, 4], SH)
+
+    # A [BM, BK], k along the columns (the wrapper makes A row-major).
+    ra = gl.arange(0, BM, layout=gl.SliceLayout(1, GA))[:, None]
+    ka = gl.arange(0, BK, layout=gl.SliceLayout(0, GA))[None, :]
+    pa = A + (m0 + ra) * stride_am + ka * stride_ak
+    ma = (m0 + ra) < M
+    # B either [BN, BK] (k contiguous: a pre-transposed weight) or [BK, BN] (n contiguous).
+    if B_K_CONTIG:
+        nb_ = gl.arange(0, BN, layout=gl.SliceLayout(1, GB))[:, None]
+        kb = gl.arange(0, BK, layout=gl.SliceLayout(0, GB))[None, :]
+    else:
+        kb = gl.arange(0, BK, layout=gl.SliceLayout(1, GB))[:, None]
+        nb_ = gl.arange(0, BN, layout=gl.SliceLayout(0, GB))[None, :]
+    pb = B + kb * stride_bk + (n0 + nb_) * stride_bn
+    mb_ = (n0 + nb_) < N
+
+    acc = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+    ta = gl.load(pa, mask=ma & (ka < K), other=0.0)
+    tb = gl.load(pb, mask=mb_ & (kb < K), other=0.0)
+    for k0 in range(0, K, BK):
+        gl.barrier()                                       # every warp is done reading the previous tile
+        if SPLIT_A:
+            inv_a = _stage_split(sa, sal, ta, True, BM, BK, HI_EXP, LO_SHIFT)
+        else:
+            stage_k_inner(sa, ta, BM, BK)
+        if SPLIT_B:
+            inv_b = _stage_split(sb, sbl, tb, B_K_CONTIG, BN, BK, HI_EXP, LO_SHIFT)
+        elif B_K_CONTIG:
+            stage_k_inner(sb, tb, BN, BK)
+        else:
+            stage_k_outer(sb, tb, BK, BN)
+        gl.barrier()
+        if k0 + BK < K:                                    # the next tile loads while the HMMAs run
+            ta = gl.load(pa + (k0 + BK) * stride_ak, mask=ma & (ka + k0 + BK < K), other=0.0)
+            tb = gl.load(pb + (k0 + BK) * stride_bk, mask=mb_ & (kb + k0 + BK < K), other=0.0)
+        if SPLIT_A or SPLIT_B:
+            t_hh = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+            t_lo = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+            for j in gl.static_range(J):
+                t_hh = mma_k4(sa, j * BM, sb, j * BN, t_hh, MB, NB, LA, LB, L)
+                if SPLIT_A:
+                    t_lo = mma_k4(sal, j * BM, sb, j * BN, t_lo, MB, NB, LA, LB, L)
+                if SPLIT_B:
+                    t_lo = mma_k4(sa, j * BM, sbl, j * BN, t_lo, MB, NB, LA, LB, L)
+            t = t_hh + t_lo * (1.0 / (1 << LO_SHIFT))
+            if SPLIT_A:
+                t = t * _bcast_rows(gl.reshape(inv_a, [MB, 2, 8]), L)
+            if SPLIT_B:
+                t = t * _bcast_cols(gl.reshape(inv_b, [NB, 2, 2, 4]), L)
+            acc = acc + t
+        else:
+            t = gl.zeros([MB, 2, 8, NB, 2, 2, 4], gl.float32, layout=L)
+            for j in gl.static_range(J):
+                t = mma_k4(sa, j * BM, sb, j * BN, t, MB, NB, LA, LB, L)
+            acc = acc + t                                  # across tiles on the FMA units, round to nearest
+
+    rows, cols = acc_rows_cols(L, MB, NB)
+    rows = m0 + rows
+    cols = n0 + cols
+    if HAS_BIAS:
+        # any broadcast of the bias: a [N] row has (0, 0, 1) strides, a full [B, M, N] its own
+        bias = gl.load(Bias + z * stride_xz + rows.to(gl.int64) * stride_xm + cols.to(gl.int64) * stride_xn,
+                       mask=(rows < M) & (cols < N), other=0.0)
+        if PROMOTE_BIAS:
+            bias = bias.to(gl.float32)
+        acc = alpha * acc + beta * bias
+    c = acc.to(C.dtype.element_ty)
+    if EPILOGUE != 0:
+        # Per-stage rounding emulation, as matmul_kernel: `c` carries the unfused GEMM's store rounding, the
+        # widening is the standalone epilogue kernel's load; the formulas are those of ops/silu.py / ops/gelu.py.
+        x = c.to(gl.float32)
+        if EPILOGUE == 1:
+            out = x * _sigmoid(x)
+        elif EPILOGUE == 2:
+            out = 0.5 * (1.0 + gl.erf(0.707106781 * x)) * x
+        else:
+            out = 0.5 * (1.0 + _tanh(0.7978845608 * x * (1.0 + 0.044715 * x * x))) * x
+        c = out.to(C.dtype.element_ty)
+    gl.store(C + rows.to(gl.int64) * stride_cm + cols.to(gl.int64) * stride_cn, c, mask=(rows < M) & (cols < N))
+
+
+@gluon.jit
+def _pow2(e):
+    """2^e as fp32 for an integer exponent in [-126, 127], built from its bits (exact)."""
+    return ((e + 127) << 23).to(gl.float32, bitcast=True)
+
+
+@gluon.jit
+def split_exponent(m, HI_EXP: gl.constexpr):
+    """The exponent s of the power-of-two scale 2^s that brings a largest magnitude m (fp32) into
+    [2^HI_EXP, 2^(HI_EXP+1)): 0 for m = 0, clamped to the fp32 normal range."""
+    e = ((m.to(gl.int32, bitcast=True) >> 23) & 0xFF) - 127
+    s = gl.where(m == 0.0, 0, HI_EXP - e)
+    return gl.minimum(gl.maximum(s, -126), 126)
+
+
+@gluon.jit
+def hi_lo(x, LO_SHIFT: gl.constexpr):
+    """fp32 x (already scaled) as hi = fp16(x) and lo = fp16((x - hi) 2^LO_SHIFT); lo is 0 where hi holds x whole
+    (an infinity among them)."""
+    hi = x.to(gl.float16)
+    h32 = hi.to(gl.float32)
+    return hi, gl.where(h32 == x, 0.0, (x - h32) * (1 << LO_SHIFT)).to(gl.float16)
+
+
+@gluon.jit
+def _stage_split(s_hi, s_lo, tile, K_INNER: gl.constexpr, R: gl.constexpr, KK: gl.constexpr,
+                 HI_EXP: gl.constexpr, LO_SHIFT: gl.constexpr):
+    """Stage one operand tile as hi + lo of the unit's operand dtype, each row (k along the columns) or column (k
+    along the rows) under its own power-of-two scale 2^s; returns 2^-s per row / column [R]. A zero row keeps
+    s = 0; s is clamped to the fp32 normal range. A value hi already holds whole (an infinity among them) has lo 0."""
+    x = tile.to(gl.float32)
+    KA: gl.constexpr = 1 if K_INNER else 0
+    s = split_exponent(gl.max(gl.abs(x), axis=KA), HI_EXP)
+    hi, lo = hi_lo(x * gl.expand_dims(_pow2(s), KA), LO_SHIFT)
+    if K_INNER:
+        stage_k_inner(s_hi, hi, R, KK)
+        stage_k_inner(s_lo, lo, R, KK)
+    else:
+        stage_k_outer(s_hi, hi, KK, R)
+        stage_k_outer(s_lo, lo, KK, R)
+    return _pow2(-s)
+
+
+@gluon.jit
+def _bcast_rows(r, L: gl.constexpr):
+    """A [MB, 2, 8] row statistic in any layout -> broadcastable over L."""
+    RL: gl.constexpr = gl.SliceLayout(3, gl.SliceLayout(4, gl.SliceLayout(5, gl.SliceLayout(6, L))))
+    return bcast_rows(gl.convert_layout(r, RL), L)
+
+
+@gluon.jit
+def _bcast_cols(c, L: gl.constexpr):
+    """A [NB, 2, 2, 4] column statistic (column = 16 o3 + 8 o4 + 4 o5 + o6, acc_rows_cols) -> broadcastable over L."""
+    CL: gl.constexpr = gl.SliceLayout(0, gl.SliceLayout(1, gl.SliceLayout(2, L)))
+    c = gl.convert_layout(c, CL)
+    return gl.expand_dims(gl.expand_dims(gl.expand_dims(c, 0), 1), 2)
+
+
+@functools.lru_cache(maxsize=None)
+def mm_layouts(MB: int, NB: int, WM: int, WN: int, BK: int, b_k_contig: bool) -> dict:
+    """Every layout the kernel takes for a tile of 16 MB x 16 NB x BK over WM x WN warps."""
+    nw = WM * WN
+    tk = min(BK // 8, 32)
+    g_rows = gl.BlockedLayout([1, 8], [32 // tk, tk], [nw, 1], [1, 0])        # [rows, k], 8 halves along k
+    return dict(
+        L=acc_layout(MB, NB, WM, WN), LA=frag_a_layout(MB, NB, WM, WN), LB=frag_b_layout(MB, NB, WM, WN),
+        GA=g_rows,
+        # [n, k] like A, or [k, n] with 4 consecutive k per thread (one STS.64 per transposing store)
+        GB=g_rows if b_k_contig else gl.BlockedLayout([4, 2], [1, 32], [nw, 1], [1, 0]))