nbxtensor 0.0.1__py3-none-any.whl

nbxtensor/ops/autotune_refusals.py

@@ -0,0 +1,380 @@
+"""A backend refusal for one candidate config costs that config, not the run.
+
+Triton's autotuner already has this semantics: `_bench` catches
+`OutOfResources`, `CompileTimeAssertionFailure` and `PTXASError`, scores that
+config `inf`, and carries on. A backend refusal says the same thing -- this
+config cannot be compiled here -- but a backend's refusal type descends from
+`RuntimeError` rather than `TritonError`, so nothing catches it and it ends
+the run.
+
+This module names NO backend and NO vendor: it asks the Metal seam
+(`triton.metal_backend`) whether an exception is a backend refusal. The seam is
+the only place an implementation is named, because this file is shared -- the
+CUDA rack runs it too, where the answer is simply False.
+
+Measured 2026-09-12: hat-s-x4 and real-esrgan-x2 both died inside the sweep on
+
+    Refusing a tt.dot with tile dim 128 (> 64) on the generic per-thread path
+
+while `conv2d_forward_kernel` declares eighteen configs of which eleven have
+every dimension at 64 or below. Eleven servable choices, and the run died on
+the first of the seven that refuse.
+
+The refusal itself is right and stays: that path is validated only to 64 and
+says so. What this changes is where it lands.
+
+Two rules, and the second is what keeps an exclusion from becoming a silent
+narrowing:
+
+  * only a REFUSAL is excluded. An out-of-memory, a kernel bug, a driver
+    fault must end the run exactly as before -- scoring those `inf` would turn
+    every failure into a shape that is merely slower.
+  * the exclusion is SAID, once per reason. A tuner that quietly drops configs
+    leaves a shape slow for a cause nobody can find.
+"""
+from __future__ import annotations
+
+from typing import Callable, Dict, Iterable, List
+
+#: Reasons already announced this process, so a sweep of eighteen configs
+#: prints one line rather than eighteen identical ones.
+_ANNOUNCED: set = set()
+
+_INF: List[float] = [float("inf"), float("inf"), float("inf")]
+
+
+def reset_announcements() -> None:
+    """Forget what has been said. For tests, which must not depend on order.
+
+    The "once per process" above is deliberate product behaviour and stays;
+    without a reset it also makes the behaviour order-dependent, so the second
+    test to use a reason sees nothing and cannot tell that from a bug.
+    """
+    _ANNOUNCED.clear()
+
+
+def _is_backend_refusal(exc: BaseException) -> bool:
+    """Is this the backend saying it cannot compile this config correctly?
+
+    Asked by class, not by message: matching on the text would make every
+    reworded refusal a run-ending error again, silently.
+
+    This SHARED module names NO backend vendor: it asks the Metal seam
+    (`triton.metal_backend.is_backend_refusal`), which knows the selectable
+    backends (the bledden fork, triton-ext). On a machine with no Metal backend
+    — the Dell on CUDA — the seam returns False, exactly as this did when it
+    imported one vendor by name. One seam decides which backend runs; the engine
+    does not branch to a vendor, here or anywhere else.
+    """
+    try:
+        from nbxtensor.backend.metal_backend import is_backend_refusal
+    except Exception:                      # the seam is not importable here
+        return False
+    return is_backend_refusal(exc)
+
+
+def exclude_refused_configs(bench: Callable, say: Callable[[str], None] | None = None):
+    """Wrap a per-config benchmark so a refusal scores `inf` instead of raising."""
+    def _say(line: str) -> None:
+        if say is not None:
+            say(line)
+        else:
+            print(line, flush=True)
+
+    def _run(*args, **kwargs):
+        try:
+            return bench(*args, **kwargs)
+        except BaseException as exc:
+            if not _is_backend_refusal(exc):
+                raise
+            reason = str(exc).strip().splitlines()[0][:160]
+            if reason not in _ANNOUNCED:
+                _ANNOUNCED.add(reason)
+                _say(f"[AUTOTUNE_REFUSED] a candidate config is not servable on "
+                     f"this backend and was excluded from the sweep: {reason}")
+            return list(_INF)
+    return _run
+
+
+def all_refused(timings: Dict) -> bool:
+    """Did EVERY config score infinite?
+
+    `min` over all-inf returns an arbitrary config, which then refuses at
+    launch with a message about that one config -- hiding that every config
+    refused. A caller that can ask this can say the true thing instead.
+    """
+    values: Iterable = timings.values() if hasattr(timings, "values") else timings
+    values = list(values)
+    if not values:
+        return False
+    return all(all(x == float("inf") for x in (v if isinstance(v, (list, tuple)) else [v]))
+               for v in values)
+
+
+# ── the per-candidate time budget ──────────────────────────────────────────
+#
+# A third species of candidate, measured 2026-09-12 on a bf16 matmul sweep at
+# M=1500: neither refused nor wrong, but PATHOLOGICALLY SLOW -- one launch of
+# 2.5 minutes, which the five-launch estimate multiplied past twelve. Neither
+# the refusal exclusion above nor the oracle sees it; only a campaign-level
+# timeout reaped it, killing the whole run without naming the candidate.
+#
+# The budget is DERIVED, never a constant:
+#   * first candidate of a sweep: R0 x the fp64 oracle's own CPU time for the
+#     key it just computed (the one number measured before any candidate) --
+#     a GPU configuration slower than the CPU float64 of the same mathematics
+#     is not a configuration;
+#   * later candidates: R x the fastest single-launch time completed so far
+#     in the same sweep (the winner sets the scale).
+#
+# The ratios themselves come from measurement, and until they are measured
+# they are None and the watchdog is DISARMED -- it must never guess. When the
+# measurement lands, its numbers are written beside the values.
+_BUDGET_RATIOS = {"first_vs_oracle_cpu": None, "later_vs_best": None}
+
+#: The most recent oracle CPU time, in ms. The screen computes the oracle for
+#: a key immediately before Triton benches that key's candidates, and sweeps
+#: are serialised in-process, so "most recent" IS this sweep's key. That
+#: serialisation is an assumption this module states rather than hides; a
+#: parallel sweep would need the key joined explicitly.
+_SWEEP: dict = {"best_ms": None, "budget_ms": None, "token": None}
+
+
+def _key_for_cache(at, args, kwargs):
+    """The key exactly as `Autotuner.run` files it in `at.cache`, so the
+    unmeasured mark lands on the same identity `capture()` will walk."""
+    try:
+        from nbxtensor.ops import autotune_cache as atc
+        return tuple(atc.key_of(at, args, kwargs))
+    except Exception:                                  # noqa: BLE001
+        return None
+
+
+def begin_sweep() -> None:
+    """Reset the per-sweep state. Called when a new key's bench begins."""
+    _SWEEP["best_ms"] = None
+    _SWEEP["budget_ms"] = None
+    _SWEEP["swap0_mb"] = None
+    _SWEEP["swap_said"] = False
+
+
+def note_sweep_swap_baseline() -> None:
+    """Snapshot swap usage at the sweep's start, from our own authority."""
+    try:
+        from nbxtensor.backend.host_memory import memory_state   # the library's own reading of the host
+        _SWEEP["swap0_mb"] = memory_state().swap_used_mb
+    except Exception:                                  # noqa: BLE001
+        _SWEEP["swap0_mb"] = None
+
+
+def sweep_grew_the_swap() -> int:
+    """MB of swap growth since the sweep's baseline, 0 when unknown or quiet.
+
+    The threshold is ZERO -- the identity, not a tuned constant -- because the
+    costs are asymmetric: a false mark (a daemon moved the swap) only skips
+    persistence and the key re-sweeps another day, while a missed mark
+    persists a choice timed against the swap, which keeps deciding on machines
+    and days it knows nothing about. Reproduced 2026-09-13: a 5.9 GB sweep
+    passed the pre-gate at 6.5 GB available, saturated the machine, and the
+    FOLLOWING key crawled in `waitUntilCompleted` with swap at 7621/8192 MB --
+    the round-9 shape, whose own arguments were 11 MB and innocent.
+    """
+    base = _SWEEP.get("swap0_mb")
+    if base is None:
+        return 0
+    try:
+        from nbxtensor.backend import device_facts as _facts, host_memory as _hm   # the library's own readings
+        host_shares_memory_with_device, memory_state = _facts.host_shares_memory_with_device, _hm.memory_state
+        if not host_shares_memory_with_device():
+            return 0                                   # discrete memory: host swap is not the bench
+        now = memory_state().swap_used_mb
+    except Exception:                                  # noqa: BLE001
+        return 0
+    if now is None:
+        return 0
+    return max(0, int(now) - int(base))
+
+
+def mark_sweep_unmeasured(at, key, grew_mb: int, say=None) -> None:
+    """Mark this sweep's choice as timed against the swap, and SAY it once."""
+    from nbxtensor.ops import autotune_cache as atc
+
+    atc.mark_unmeasured(at, key)
+    if not _SWEEP.get("swap_said"):
+        _SWEEP["swap_said"] = True
+        line = (f"[AUTOTUNE_SWAP] the swap grew by {grew_mb} MB during this "
+                f"sweep: its timings measured the swap, not the kernels. The "
+                f"choice stands for this run and will NOT be persisted.")
+        (say or (lambda l: print(l, flush=True)))(line)
+
+
+def current_budget_ms():
+    """The budget the bench should enforce for the CURRENT candidate, or None.
+
+    Read by the launcher's `do_bench` when its caller passed none: Triton's
+    `_bench -> self.do_bench(kernel_call, quantiles=...)` chain is not ours to
+    re-sign, so the wrapper publishes here and the bench consults it. Explicit
+    coupling, documented at both ends.
+    """
+    return _SWEEP.get("budget_ms")
+
+
+def _compute_budget():
+    later = _BUDGET_RATIOS["later_vs_best"]
+    first = _BUDGET_RATIOS["first_vs_oracle_cpu"]
+    if _SWEEP["best_ms"] is not None:
+        return None if later is None else later * _SWEEP["best_ms"]
+    if first is None:
+        return None
+    try:
+        from nbxtensor.ops.screen_oracle import _ORACLE_MS
+        if _ORACLE_MS:
+            return first * next(reversed(_ORACLE_MS.values()))
+    except Exception:                                  # noqa: BLE001
+        pass
+    return None
+
+
+def note_candidate_time(single_launch_ms: float) -> None:
+    """Feed a completed candidate's probe time back into the sweep state."""
+    best = _SWEEP.get("best_ms")
+    if best is None or single_launch_ms < best:
+        _SWEEP["best_ms"] = float(single_launch_ms)
+
+
+def exclude_slow_candidates(bench, say=None):
+    """Wrap a per-config benchmark so an over-budget candidate scores `inf`.
+
+    Same shape as `exclude_refused_configs`, same two rules: only the
+    dedicated exception is caught -- everything else propagates -- and the
+    exclusion is SAID with both of its numbers, because a candidate scored
+    out for time without its time is a silent narrowing.
+    """
+    def _say(line):
+        (say or (lambda l: print(l, flush=True)))(line)
+
+    def _run(*args, **kwargs):
+        from nbxtensor.backend.launcher import CandidateOverTimeBudget
+
+        _SWEEP["budget_ms"] = _compute_budget()
+        try:
+            out = bench(*args, **kwargs)
+        except CandidateOverTimeBudget as exc:
+            _say(f"[AUTOTUNE_SLOW] a candidate took {exc.took_ms:.1f} ms for "
+                 f"ONE launch against a budget of {exc.budget_ms:.1f} ms "
+                 f"derived from this sweep's own measurements; excluded from "
+                 f"the sweep, which continues")
+            return list(_INF)
+        finally:
+            _SWEEP["budget_ms"] = None
+        return out
+    return _run
+
+
+# ── the per-config memory footprint door ───────────────────────────────────
+#
+# `bench_would_swap` (launcher) refuses to bench a KEY whose ARGUMENTS alone
+# exceed available memory. But a sweep tries CONFIGS, and a config allocates
+# more than the arguments: Triton's L2-flush bench scratch (a 256 MiB torch
+# buffer per never-seen key) plus the compiled kernel's own global/profile
+# scratch. A config whose full footprint exceeds what remains is caught only
+# AFTER it has run and grown the swap. This door estimates the footprint and
+# excludes such a config BEFORE the spend -- the threadgroup-budget door
+# (codegen refuses an over-budget tile) transposed to unified memory.
+#
+# The limit is read from `core.host_memory` -- the SAME source Prism plans
+# against, never `Pages free` -- and live, because pressure is a state of the
+# moment. The estimate is a LOWER bound made honest by a margin: undercounting
+# would let a swapping config through, so unknown scratch is treated as its
+# declared size and a safety margin is added, never subtracted.
+
+#: Triton's L2-flush bench buffer, allocated once per never-seen key.
+_BENCH_SCRATCH_BYTES = 256 * 1024 * 1024
+#: A fraction of available kept free so the estimate's own slack cannot tip
+#: the machine over. Not a tuned magic number -- a headroom, stated.
+_FOOTPRINT_MARGIN = 0.10
+
+
+def config_footprint_fits(arg_bytes: int, config_scratch_bytes: int = 0,
+                          available_mb: "float | None" = None):
+    """(fits, available_mb, footprint_mb) for one candidate config.
+
+    footprint = arguments + the kernel's own scratch + the bench L2 buffer.
+    Compared to available memory less a margin. `available_mb` defaults to the
+    live `core.host_memory` reading (Prism's source). An unreadable platform
+    (None) fits everything and says nothing -- it adds no policy of its own.
+    """
+    if available_mb is None:
+        try:
+            from nbxtensor.backend import host_memory as _hm   # the library's own reading of the host
+            memory_state = _hm.memory_state
+            available_mb = memory_state().available_mb
+        except Exception:                              # noqa: BLE001
+            available_mb = None
+    if available_mb is None:
+        return True, None, None
+    footprint = int(arg_bytes) + int(config_scratch_bytes) + _BENCH_SCRATCH_BYTES
+    usable = available_mb * (1.0 - _FOOTPRINT_MARGIN) * 2 ** 20
+    return footprint <= usable, available_mb, footprint / 2 ** 20
+
+
+def exclude_over_footprint(bench, arg_bytes, config_scratch_bytes=0, say=None):
+    """Wrap a per-config bench so an over-footprint config scores `inf` without
+    running -- excluded BEFORE the spend, and SAID with its numbers."""
+    def _say(line):
+        (say or (lambda l: print(l, flush=True)))(line)
+
+    def _run(*args, **kwargs):
+        fits, avail, fp_mb = config_footprint_fits(arg_bytes, config_scratch_bytes)
+        if not fits:
+            _say(f"[AUTOTUNE_FOOTPRINT] a candidate config needs ~{fp_mb:.0f} MB "
+                 f"(args + scratch + bench buffer) against {avail:.0f} MB "
+                 f"available; excluded BEFORE running it, so the sweep neither "
+                 f"swaps nor measures the swap. The sweep continues.")
+            return list(_INF)
+        return bench(*args, **kwargs)
+    return _run
+
+
+def install() -> bool:
+    """Wrap `Autotuner._bench` so a refused config is excluded, once.
+
+    Installed by the launcher beside the screen oracle. Returns False when the
+    autotuner is not importable, so a caller can say so rather than assume it
+    took.
+    """
+    try:
+        from triton.runtime.autotuner import Autotuner
+    except Exception:
+        return False
+    if getattr(Autotuner._bench, "_nbx_excludes_refusals", False):
+        return True
+    original = Autotuner._bench
+
+    def _bench(self, *args, config=None, **kwargs):
+        # A sweep boundary is a change of (tuner, live-arg key): Triton
+        # benches one key's candidates consecutively, so the first _bench of
+        # a new pair is the first candidate of a new sweep. Detected HERE
+        # because nothing upstream of _bench is ours to hook -- and without
+        # this call, `begin_sweep` would be machinery with no caller, the
+        # register's own vacuous form.
+        try:
+            from nbxtensor.ops import autotune_cache as _atc
+            self.nargs = dict(zip(self.arg_names, args))
+            _token = (id(self), str(_atc.key_of(self, args, kwargs)))
+        except Exception:                              # noqa: BLE001
+            _token = (id(self), None)
+        if _SWEEP.get("token") != _token:
+            begin_sweep()
+            _SWEEP["token"] = _token
+            note_sweep_swap_baseline()
+        _out = exclude_slow_candidates(exclude_refused_configs(
+            lambda: original(self, *args, config=config, **kwargs)))()
+        _grew = sweep_grew_the_swap()
+        if _grew:
+            mark_sweep_unmeasured(self, _key_for_cache(self, args, kwargs), _grew)
+        return _out
+
+    _bench._nbx_excludes_refusals = True
+    Autotuner._bench = _bench
+    return True

nbxtensor/ops/autotune_settings.py

@@ -0,0 +1,58 @@
+"""The autotune's two settings an application may set, in a module that loads nothing but the standard library — so
+an application can set them when it starts without importing the autotune itself.
+
+- `set_store(path)`: the directory of proven kernel settings (a path, or a callable returning one) (None: the library's own, the user's cache,
+  `backend.device_facts.certified_dir`); `directory()` the one in force.
+- `set_miss_explainer(fn)`: `fn(kernel_qual, key, memory_gb) -> str`, an application's account of a key a confirmation
+  run misses (where it expected that key to come from), appended to the refusal; None removes it.
+
+`ops.autotune_certified` reads both at call time.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Callable, Dict, Optional
+
+__all__ = ["set_store", "directory", "set_miss_explainer", "miss_explainer"]
+
+_STORE: Dict[str, object] = {}        # {"dir": a Path, or a callable returning one} once a caller set its store
+_MISS_EXPLAINER: Optional[Callable] = None
+
+
+def set_store(path) -> None:
+    """Set the directory of proven kernel settings: a path, or a callable returning one (read each time the store is
+    asked — an application can set it when it starts without resolving it then); `None` restores the library's own,
+    the user's cache (`backend.device_facts.certified_dir`)."""
+    if path is None:
+        _STORE.clear()
+    else:
+        _STORE["dir"] = path if callable(path) else Path(path)
+
+
+def directory() -> Path:
+    """The store in force: NEUROBRIX_AUTOTUNE_CERTIFIED_DIR when set (tests, a contributor's draft), else the one a
+    caller set (`set_store`), else the library's own in the user's cache."""
+    override = os.environ.get("NEUROBRIX_AUTOTUNE_CERTIFIED_DIR")
+    if override:
+        return Path(override)
+    if "dir" in _STORE:
+        where = _STORE["dir"]
+        return Path(where()) if callable(where) else where
+    from nbxtensor.backend.device_facts import certified_dir
+    return certified_dir()
+
+
+def set_miss_explainer(explain) -> None:
+    """Install `explain(kernel_qual, key, memory_gb) -> str`: where the application expected a key it finds uncertified
+    to come from, appended to the refusal; None removes it. Optional: without one, the refusal says what the store holds
+    and nothing more."""
+    global _MISS_EXPLAINER
+    if explain is not None and not callable(explain):
+        raise TypeError("set_miss_explainer: the explainer must be callable (kernel_qual, key, memory_gb) -> str")
+    _MISS_EXPLAINER = explain
+
+
+def miss_explainer() -> Optional[Callable]:
+    """The explainer in force, or None."""
+    return _MISS_EXPLAINER

nbxtensor/ops/avg_pool2d.py

@@ -0,0 +1,73 @@
+"""Average pooling 2D — pure @triton.jit kernel.
+
+Extracted from FlagGems avg_pool2d.py. Forward only (inference).
+2D blocked approach with proper padding, dilation, count_include_pad support.
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def avg_pool2d_forward_kernel(
+    input_ptr,
+    output_ptr,
+    in_stride_n, in_stride_c, in_stride_h, in_stride_w,
+    in_c, in_h, in_w,
+    out_h, out_w,
+    kernel_h: tl.constexpr,
+    kernel_w: tl.constexpr,
+    stride_h: tl.constexpr,
+    stride_w: tl.constexpr,
+    padding_h: tl.constexpr,
+    padding_w: tl.constexpr,
+    dilation_h: tl.constexpr,
+    dilation_w: tl.constexpr,
+    COUNT_INCLUDE_PAD: tl.constexpr,
+    divisor_override,
+    BLOCK_H: tl.constexpr,
+    BLOCK_W: tl.constexpr,
+):
+    pid_nc = tl.program_id(0).to(tl.int64)
+    pid_hw = tl.program_id(1).to(tl.int64)
+    num_w_blocks = tl.cdiv(out_w, BLOCK_W)
+    h_block_idx = pid_hw // num_w_blocks
+    w_block_idx = pid_hw % num_w_blocks
+    n_idx = pid_nc // in_c
+    c_idx = pid_nc % in_c
+
+    h_out_offsets = h_block_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    w_out_offsets = w_block_idx * BLOCK_W + tl.arange(0, BLOCK_W)
+
+    sum_acc = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.float32)
+    count_acc = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32)
+
+    input_base_ptr = input_ptr + n_idx * in_stride_n + c_idx * in_stride_c
+
+    for kh in range(0, kernel_h):
+        for kw in range(0, kernel_w):
+            h_in = h_out_offsets[:, None] * stride_h - padding_h + kh * dilation_h
+            w_in = w_out_offsets[None, :] * stride_w - padding_w + kw * dilation_w
+            in_mask = (h_in >= 0) & (h_in < in_h) & (w_in >= 0) & (w_in < in_w)
+
+            input_offset = h_in * in_stride_h + w_in * in_stride_w
+            current_val = tl.load(input_base_ptr + input_offset, mask=in_mask, other=0.0)
+
+            sum_acc += tl.where(in_mask, current_val, 0.0)
+            count_acc += in_mask.to(tl.int32)
+
+    if divisor_override != 0:
+        divisor = tl.full((BLOCK_H, BLOCK_W), divisor_override, dtype=tl.float32)
+    elif COUNT_INCLUDE_PAD:
+        divisor = tl.full((BLOCK_H, BLOCK_W), kernel_h * kernel_w, dtype=tl.float32)
+    else:
+        divisor = count_acc.to(tl.float32)
+
+    output_vals = tl.where(divisor != 0, sum_acc / divisor, 0.0)
+
+    out_base_ptr = output_ptr + pid_nc * out_h * out_w
+    output_block_ptr = (
+        out_base_ptr + h_out_offsets[:, None] * out_w + w_out_offsets[None, :]
+    )
+    out_mask = (h_out_offsets[:, None] < out_h) & (w_out_offsets[None, :] < out_w)
+    tl.store(output_block_ptr, output_vals.to(output_ptr.type.element_ty), mask=out_mask)

nbxtensor/ops/baddbmm_op.py

@@ -0,0 +1,132 @@
+"""Batched matmul with optional bias (baddbmm/bmm) — pure @triton.jit kernel.
+
+Originally ported from FlagGems baddbmm (Apache-2.0 license); body aligned
+on kernels/ops/matmul.py:matmul_kernel (P-KERNEL-LAUNCH-HYGIENE B1,
+2026-07) so the SAME batched kernel serves both:
+
+  - aten::baddbmm  (HAS_BIAS=True):  out = beta * bias + alpha * (b1 @ b2)
+  - aten::bmm      (HAS_BIAS=False): out = A @ B, one 3D-grid launch for the
+    whole batch instead of the former per-batch (and per-row in decode)
+    Python launch loops in wrappers.bmm.
+
+batch1 is [B, M, K], batch2 is [B, K, N], bias broadcastable to [B, M, N].
+Grid: (cdiv(M, BLOCK_M) * cdiv(N, BLOCK_N), 1, B). Accumulates in fp32.
+
+IEEE_PRECISION / PROMOTE_B carry the exact matmul_kernel semantics (see its
+docstring): ieee dot for fp32-promoted inputs on pre-Ampere, in-register
+fp16→fp32 b-tile promotion for the mixed activation/weight case. The K-loop
+body (3-arg fused `tl.dot`, %-wrap row/col offsets, K-tail masks) is a
+line-for-line mirror of matmul_kernel so a batched launch is value-identical
+to the former per-batch matmul_kernel loop under the same autotune config.
+
+Phase 1.5 (2026-05): @triton.autotune ENABLED with the same Volta-viable
+configs as kernels/ops/matmul.py. See CLAUDE.md "Autotune policy" section
+for the doctrinal exception.
+"""
+
+import triton
+
+from nbxtensor.ops._configs import nbx_autotune
+import triton.language as tl
+
+from nbxtensor.ops.matmul import _MATMUL_AUTOTUNE_CONFIGS
+
+
+@nbx_autotune(configs=_MATMUL_AUTOTUNE_CONFIGS,
+                 key=['M_BUCKET', 'N_BUCKET', 'K_BUCKET', 'IEEE_PRECISION', 'PROMOTE_B',
+                      'HAS_BIAS'],
+                 cache_results=True)
+@triton.jit
+def baddbmm_kernel(
+    A_ptr,
+    B_ptr,
+    out_ptr,
+    bias_ptr,
+    alpha,
+    beta,
+    M, N, K, M_BUCKET, N_BUCKET, K_BUCKET,
+    stride_ab, stride_am, stride_ak,
+    stride_bb, stride_bk, stride_bn,
+    stride_ob, stride_om, stride_on,
+    bias_batch_stride,
+    bias_m_stride,
+    bias_n_stride,
+    IEEE_PRECISION: tl.constexpr = False,
+    PROMOTE_B: tl.constexpr = False,
+    HAS_BIAS: tl.constexpr = True,
+    BLOCK_M: tl.constexpr = 64,
+    BLOCK_N: tl.constexpr = 64,
+    BLOCK_K: tl.constexpr = 32,
+    GROUP_M: tl.constexpr = 8,
+):
+    """out[b] = beta*bias[b] + alpha*(A[b] @ B[b])   (HAS_BIAS=True)
+       out[b] = A[b] @ B[b]                          (HAS_BIAS=False)
+
+    Grid: (cdiv(M, BLOCK_M) * cdiv(N, BLOCK_N), 1, batch)
+    Accumulates in fp32 for numerical stability; in-register cast to the
+    output pointer's dtype before store (matmul_kernel pattern).
+    """
+    # Batch offset
+    pid_b = tl.program_id(2).to(tl.int64)      # a batch offset times a stride wraps past 2^31 elements in int32
+    A_ptr += pid_b * stride_ab
+    B_ptr += pid_b * stride_bb
+    out_ptr += pid_b * stride_ob
+    if HAS_BIAS:
+        bias_ptr += pid_b * bias_batch_stride
+
+    # tl.assume integer-analyzer hints (matmul_kernel pattern). Batch and
+    # bias strides are excluded: they are legitimately 0 for broadcast.
+    tl.assume(stride_am > 0); tl.assume(stride_ak > 0)
+    tl.assume(stride_bk > 0); tl.assume(stride_bn > 0)
+    tl.assume(stride_om > 0); tl.assume(stride_on > 0)
+
+    # 2D tile indexing with grouping for L2 locality
+    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)
+    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)
+
+    # Accumulate matmul in fp32 — line-for-line matmul_kernel K loop.
+    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_B:
+            b = b.to(a.dtype)
+        # 3-arg HMMA-FMA fused form (tutorial pattern) — same as
+        # matmul_kernel so the bmm path stays bit-identical to the former
+        # per-batch loop. IEEE_PRECISION supersedes the historical
+        # `allow_tf32=False` (equivalent lowering for fp32 inputs).
+        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
+
+    offs_cm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_cn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    out_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+
+    if HAS_BIAS:
+        # Load bias and apply alpha/beta
+        bias_ptrs = (bias_ptr + offs_cm[:, None].to(tl.int64) * bias_m_stride
+                     + offs_cn[None, :].to(tl.int64) * bias_n_stride)
+        bi = tl.load(bias_ptrs, mask=out_mask, other=0.0)
+        accumulator = accumulator * alpha + bi * beta
+
+    # In-kernel cast accum → output dtype (matmul_kernel pattern).
+    c = accumulator.to(out_ptr.dtype.element_ty)
+    o_ptrs = out_ptr + offs_cm[:, None].to(tl.int64) * stride_om + offs_cn[None, :].to(tl.int64) * stride_on
+    tl.store(o_ptrs, c, mask=out_mask)