Kernels:

kernels-community
/

finegrained-kernels

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on 2 days ago

Commit

74a1379

verified ·

1 Parent(s): 03b2bb8

Uploaded using `kernel-builder`.

Browse files

Files changed (12) hide show

README.md +1 -0
build/torch-cuda/__init__.py +3 -0
build/torch-cuda/_ops.py +1 -1
build/torch-cuda/autotune_config.py +70 -0
build/torch-cuda/bayesian_autotuner.py +138 -93
build/torch-cuda/compat.py +12 -0
build/torch-cuda/embedding.py +306 -0
build/torch-cuda/metadata.json +11 -9
build/torch-cuda/metadata.json.sigstore +1 -1
build/torch-cuda/pruners.py +1 -1
build/torch-cuda/quant.py +7 -9
build/torch-cuda/symbols.json +68 -14

README.md CHANGED Viewed

@@ -23,6 +23,7 @@ matmul_2d(...)
 - `matmul_grouped`
 - `moe_fused_batched`
 - `moe_fused_grouped`
 - `backward`
 - `get_supported_act_fns`
 - `get_supported_norms`

 - `matmul_grouped`
 - `moe_fused_batched`
 - `moe_fused_grouped`
+- `embedding_lookup`
 - `backward`
 - `get_supported_act_fns`
 - `get_supported_norms`

build/torch-cuda/__init__.py CHANGED Viewed

@@ -17,6 +17,7 @@ from .matmul import matmul_2d
 from .batched import matmul_batched
 from .grouped import matmul_grouped
 from .moe import moe_fused_batched, moe_fused_grouped
 from . import backward
 from .formats import get_supported_act_fns
 from .norm import get_supported_norms, rms_norm_rows
@@ -32,6 +33,8 @@ __all__ = [
     # the fused MoE forwards over them
     "moe_fused_batched",
     "moe_fused_grouped",
     # importing it registers the dgrad formulas on the ops, so a forward call differentiates
     "backward",
     "get_supported_act_fns",

 from .batched import matmul_batched
 from .grouped import matmul_grouped
 from .moe import moe_fused_batched, moe_fused_grouped
+from .embedding import embedding_lookup
 from . import backward
 from .formats import get_supported_act_fns
 from .norm import get_supported_norms, rms_norm_rows
     # the fused MoE forwards over them
     "moe_fused_batched",
     "moe_fused_grouped",
+    # selected-row embedding lookup
+    "embedding_lookup",
     # importing it registers the dgrad formulas on the ops, so a forward call differentiates
     "backward",
     "get_supported_act_fns",

build/torch-cuda/_ops.py CHANGED Viewed

@@ -25,7 +25,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "finegrained_kernels"
-    unique_id = "43f1581"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

 def _find_ops_name() -> str:
     kernel_name = "finegrained_kernels"
+    unique_id = "ec5cc90"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

build/torch-cuda/autotune_config.py ADDED Viewed

	@@ -0,0 +1,70 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""``FINEGRAINED_AUTOTUNE_CONFIG``: the one runtime control over every tuned kernel.
+Its value is a mode name (``fast``), inline JSON, or a path to a JSON file. JSON keys are the
+tuner's own options; ``"mode"`` there applies a preset whose keys explicit ones override.
+The tuner reads this when each kernel is decorated, i.e. when the package is imported.
+"""
+import json
+import os
+from pathlib import Path
+ENV_VAR = "FINEGRAINED_AUTOTUNE_CONFIG"
+# Tuner kwargs by name, plus the debugging outputs (``key_log_path``, ``call_count_path``,
+# ``tracebacks``) and the cross-rank tune lock wait.
+KEYS = {
+    "n_trials": int, "strict": bool, "runoff": bool, "basin_probes": bool,
+    "max_refine_iters": int, "log_path": str,
+    "key_log_path": str, "call_count_path": str, "tracebacks": bool, "lock_timeout": float,
+}
+MODES = {
+    # one measured config, no search phases: correctness runs, not performance
+    "fast": {"n_trials": 1, "runoff": False, "basin_probes": False, "max_refine_iters": 0},
+}
+def _read(value: str) -> dict:
+    """The raw config, a mode name kept as ``{"mode": ...}`` so explicit keys compose with it."""
+    if not value:
+        return {}
+    if value.lstrip().startswith("{"):
+        return json.loads(value)
+    if os.path.isfile(value):
+        return json.loads(Path(value).read_text())
+    return {"mode": value}
+def load_autotune_config() -> dict:
+    """Validated overrides from the environment: the mode's preset, then explicit keys."""
+    config = _read(os.environ.get(ENV_VAR, ""))
+    mode = config.pop("mode", None)
+    if mode is not None:
+        if mode not in MODES:
+            raise ValueError(
+                f"{ENV_VAR} mode {mode!r} is not one of {sorted(MODES)} (or a JSON object / JSON file path)"
+            )
+        config = {**MODES[mode], **config}
+    unknown = set(config) - set(KEYS)
+    if unknown:
+        raise ValueError(
+            f"{ENV_VAR} has unknown keys {sorted(unknown)}; expected a subset of {sorted(KEYS)} or \"mode\""
+        )
+    return {key: KEYS[key](value) for key, value in config.items()}

build/torch-cuda/bayesian_autotuner.py CHANGED Viewed

@@ -37,6 +37,7 @@ import random
 import time
 from collections import defaultdict
 from contextlib import contextmanager
 import torch
@@ -53,19 +54,24 @@ from triton.runtime.autotuner import (
     triton_key,
 )
 logger = logging.getLogger(__name__)
-# Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=<path>): every
-# BayesianAutotuner.run() call increments its kernel's tally, dumped to <path> at exit.
 # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token
 # re-execution — the decode-gap forensics tool.
 _CALL_COUNTS = None
-if os.environ.get("FINEGRAINED_AUTOTUNE_CALLCOUNT"):
     import atexit
     _CALL_COUNTS = {}
-    def _dump_call_counts(path=os.environ["FINEGRAINED_AUTOTUNE_CALLCOUNT"]):
         with open(path, "w") as f:
             json.dump(_CALL_COUNTS, f, indent=1, sort_keys=True)
@@ -84,48 +90,58 @@ class BayesianAutotuner(Autotuner):
         max_failures: int | None = None,
         n_startup_trials: int = 12,
         gamma: float = 0.25,
-        refine: bool = True,
         max_refine_iters: int = 5,
         log_path: str | None = None,
         path_anchor_axes: tuple[str, ...] = (),
         finite_check_args: tuple[str, ...] = (),
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.fn_name = getattr(self.fn, "__name__", str(self.fn))
-        # Bayesian trial budget — the per-decorator default, overridable via the
-        # FINEGRAINED_AUTOTUNE_TRIALS env var (quick sweeps / exhaustive runs without touching
-        # the decorators; set it >= grid size to fall back to stock exhaustive bench-all).
-        self.n_trials = int(os.environ.get("FINEGRAINED_AUTOTUNE_TRIALS") or n_trials)
         # A tune aborts once this many configs have failed to compile/run without the measured
         # budget being met. Defaults to n_trials: a grid that cannot land its measurements without
         # that many rejects has a pruner gap, and the fix is to fence the dead region rather than
         # to tolerate the compiles. Its own axis so a kernel can raise it deliberately.
-        # an explicit argument wins over the env override, which wins over the default; `or`
-        # would drop a deliberate `max_failures=0`, the strictest setting there is
-        if max_failures is None:
-            env_max_failures = os.environ.get("FINEGRAINED_AUTOTUNE_MAX_FAILURES")
-            max_failures = env_max_failures if env_max_failures is not None else self.n_trials
         self.max_failures = int(max_failures)
         self.n_startup_trials = n_startup_trials
         # top fraction of measured configs the TPE treats as "good"
         self.gamma = gamma
-        self.refine = refine
-        self.max_refine_iters = max_refine_iters
         # Branch axes (by kwarg name) that relocate the TILE optimum — the caller's
         # declaration, since it is configuration knowledge (the tuner never hardcodes
         # axis names). Each distinct value combination of these axes gets one guaranteed
         # max-tile anchor; empty = median anchors only.
         self.path_anchor_axes = tuple(path_anchor_axes)
         # Output arg NAMES (a declaration — the tuner never hardcodes kernel details) whose
         # values are checked for NaN/inf after each config's bench: a config that writes
         # non-finite output scores inf and can never be crowned. The tuner times, it never
         # validates — this veto closes the NaN half of the "wrong-answer configs win
         # tunes" class at negligible cost (one isfinite pass per benched config).
         self.finite_check_args = tuple(finite_check_args)
-        # JSONL log of every benched (key, config, ms) — set here or via the
-        # FINEGRAINED_AUTOTUNE_LOG env var. Analyse offline to prune bad configs.
-        self.log_path = log_path or os.environ.get("FINEGRAINED_AUTOTUNE_LOG")
         # steady-state launch support: key extraction without dict builds (eager decode
         # is host-bound; the cache-hit path of run() is on the per-token critical path)
         self._arg_position = {name: i for i, name in enumerate(self.arg_names)}
@@ -195,7 +211,7 @@ class BayesianAutotuner(Autotuner):
             if root is not e:
                 msg += f" || root: {type(root).__name__}: {str(root)[-300:]}"
             err = f"{type(e).__name__}: {msg}"
-            if os.environ.get("FINEGRAINED_AUTOTUNE_TRACEBACKS"):
                 logger.warning(
                     "[autotune] %s failed on %s:\n%s",
                     self.fn_name,
@@ -460,9 +476,22 @@ class BayesianAutotuner(Autotuner):
         self.nargs = None
         return ret
     def _bayesian_search(self, configs: List[Config], args, kwargs, key) -> Config:
         timings: Dict[int, float] = {}
-        sigs = [tuple(sorted(c.all_kwargs().items())) for c in configs]
         def bench_idx(idx: int) -> float:
             if idx in timings:
@@ -489,8 +518,8 @@ class BayesianAutotuner(Autotuner):
         # per-dimension model can write off a whole axis it never saw succeed (two dsv4 tunes
         # shipped 25-60% slow winners because their random seeds only sampled a basin's dead
         # configs), and coordinate descent can't recover a winner two coupled flips away.
-        n_startup = max(2, min(self.n_startup_trials, self.n_trials))
-        anchors = self._basin_anchor_indices(configs)
         order = list(range(len(configs)))
         random.Random(0).shuffle(order)
         warm_idx = self._warm_start_index(configs, key)
@@ -574,22 +603,21 @@ class BayesianAutotuner(Autotuner):
         # Coordinate-descent refinement: try single-dim perturbations around
         # the current best until no neighbor improves.
-        if self.refine:
-            for _ in range(self.max_refine_iters):
-                best_idx = min(timings, key=timings.get)
-                best_sig = sigs[best_idx]
-                best_ms = timings[best_idx]
-                improved = False
-                for i, s in enumerate(sigs):
-                    if i in timings:
-                        continue
-                    diff = sum(1 for (_, a), (_, b) in zip(best_sig, s) if a != b)
-                    if diff != 1:
-                        continue
-                    if bench_idx(i) < best_ms:
-                        improved = True
-                if not improved:
-                    break
         # Crown RUNOFF: the search's single-shot timings carry enough noise to flip crowns
         # between same-code tune rolls (gpt-oss prefill grouped rolled 1727-2150us on identical
@@ -597,21 +625,24 @@ class BayesianAutotuner(Autotuner):
         # the top candidates back-to-back — same warmed process, adjacent in time, so they see
         # the same clocks — and crown the winner of THAT comparison. Costs ~3 extra benches.
         finite = sorted((t, i) for i, t in timings.items() if t != float("inf"))
-        if len(finite) > 1:
-            for _, i in finite[:3]:
                 timings.pop(i, None)
                 bench_idx(i)  # memoizing closure: re-benches and re-records the popped entry
         self.configs_timings = {configs[i]: t for i, t in timings.items()}
-        best = min(timings, key=timings.get)
         if timings[best] == float("inf"):
             # crowning an arbitrary broken config would persist it as `best` on disk and
             # every later healthy process would replay it without re-tuning — the exact
             # sticky-contaminated-context scenario inf-forgiveness exists to survive
             self._report_bench_failures()
             raise RuntimeError(
-                f"[bayesian-autotune] {self.fn_name}: every benched config failed "
-                "(all-inf tune) — nothing correct to crown or persist; the context may "
                 "be trap-contaminated (retry in a fresh process) or the grid is broken "
                 "for this launch (see the failure report above)."
             )
@@ -639,12 +670,8 @@ class BayesianAutotuner(Autotuner):
             pass
     def _basin_anchor_indices(self, configs: List[Config]) -> List[int]:
-        """One representative config index per (COMPUTE_MODE, SWAP_AB) basin — the MEDIAN in
-        tile-sort order, a mid-sized tile with mid warps/stages. Not the smallest: a basin's
-        minimal corner (min BK x 2 warps) can be latency-bound pathological (a 131µs anchor in
-        a basin whose peak is 41µs re-poisons the axis it was meant to protect). Coordinate
-        descent climbs BN/BK/warps/stages from wherever the TPE lands within the basin.
-        Returns [] when the grid has no such axes (single basin)."""
         # Basin axes are DERIVED, not declared: a config kwarg with string or boolean
         # values is a branch axis (different code path — compute mode, operand swap,
         # warp specialization), which partitions the grid into disjoint performance
@@ -670,41 +697,46 @@ class BayesianAutotuner(Autotuner):
         if len(groups) <= 1:
             return []
-        def tile_order(i):
-            return (
-                configs[i].kwargs.get("BLOCK_SIZE_N", 0),
-                configs[i].kwargs.get("BLOCK_SIZE_K", 0),
-                configs[i].num_warps,
-                configs[i].num_stages,
-            )
-        # Median anchor per basin, PLUS one max-tile anchor per declared PATH-axis combo
-        # (``path_anchor_axes``). The median alone left the big-tile end unsampled: TPE
-        # densities lock onto the first-timed mid-tile region and coordinate descent cannot
-        # cross DIAGONAL ridges — dsv4 mxfp8 prefill gate_up's BN=128/BK=256 winner (840us)
-        # needs BN and warps to move together, so descent from the BN=64 crown rejects the
-        # lone BN step and n_trials 100/200/400 all missed it; one guaranteed big-tile
-        # sample fixed it (crown 958 -> 880 at the default budget). The max anchor rides
-        # the declared COARSE grouping, NOT the full branch-axis cross: per-basin max
-        # anchors on the weight-only grouped grid (32 basins: modes x WS x PS x 2 memory
-        # modes each way) put 64 forced trials into a 100-trial budget and starved the
-        # search — gpt-oss prefill regressed 3463 -> 4099us on a quiet box before this was
-        # scoped. Which axes qualify is the KERNEL's declaration (configuration knowledge);
-        # axes absent from this grid are ignored.
         anchors = []
         coarse_seen = set()
         coarse_axes = [k for k in self.path_anchor_axes if k in basin_axes]
         for basin, idxs in groups.items():
-            ordered = sorted(idxs, key=tile_order)
-            anchors.append(ordered[len(ordered) // 2])
             if not self.path_anchor_axes:
-                continue  # no declaration -> median anchors only
-            # declared axes absent from THIS grid still coarse-group to () — one global
-            # max-tile anchor, the big-tile guarantee without per-basin anchor spam
             coarse = tuple(basin[basin_axes.index(k)] for k in coarse_axes)
-            if coarse not in coarse_seen and ordered[-1] != anchors[-1]:
                 coarse_seen.add(coarse)
-                anchors.append(ordered[-1])
         return anchors
     def _warm_start_index(self, configs: List[Config], key=None):
@@ -729,8 +761,8 @@ class BayesianAutotuner(Autotuner):
         """Persist the tuned best config to Triton's on-disk cache so a later run (or process)
         skips the whole minutes-long search+compile — most of which is Triton codegen, not
         benching. Keyed exactly like Triton's own autotune cache (triton version + backend +
-        kernel source hash + invalidating env + key + the full config grid), so any of those
-        changing re-tunes.
         Unlike Triton's stock version we do NOT bail when configs carry a ``pre_hook`` (ours
         always do): we persist each config's ``all_kwargs()`` — serializable and pre_hook-free —
@@ -750,19 +782,22 @@ class BayesianAutotuner(Autotuner):
             triton_key(),
             make_backend(driver.active.get_current_target()).hash(),
             fn.cache_key,
             str(sorted(env_vars.items())),
             # Key NAMES and the search budget are determinants too: tuning_key holds
             # values only (renaming a key or swapping two same-typed keys would collide),
-            # and FINEGRAINED_AUTOTUNE_TRIALS is not in Triton's invalidating-env list —
             # without n_trials a 3-trial exploratory winner gets replayed at 100 trials.
             str(self.keys),
-            f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}",
             str(tuning_key),
         ] + [str(c) for c in configs]  # str(Config) is pre_hook-free and process-stable
         cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest()
         cache = get_cache_manager(cache_key)
         file_name = f"{fn.__name__[:150]}.bayes_autotune.json"
-        keylog = os.environ.get("FINEGRAINED_AUTOTUNE_KEYLOG")
         if keylog:
             with open(keylog, "a") as f:
                 f.write(json.dumps({
@@ -840,7 +875,7 @@ def _tuning_lock(cache, file_name: str):
     Held while a key is tuned and its crown published, so concurrent TP/EP ranks converge on the
     same config instead of each crowning its own (see ``check_disk_cache``). Advisory ``flock`` on
     a sidecar file next to the crown; entirely best-effort — no cache dir, no ``fcntl``, or a lock
-    that cannot be taken within ``FINEGRAINED_TUNE_LOCK_TIMEOUT`` (default 30 min, since the holder
     is doing a real tune) all fall through to tuning locally. Losing the lock costs redundant work,
     never correctness."""
     lock_file = None
@@ -850,9 +885,7 @@ def _tuning_lock(cache, file_name: str):
         cache_dir = getattr(cache, "cache_dir", None)
         if cache_dir:
             os.makedirs(cache_dir, exist_ok=True)
-            deadline = time.time() + float(
-                os.environ.get("FINEGRAINED_TUNE_LOCK_TIMEOUT") or 1800
-            )
             lock_file = open(os.path.join(cache_dir, f"{file_name}.lock"), "w")
             while True:
                 try:
@@ -892,10 +925,12 @@ def bayesian_autotune(
     max_failures: int | None = None,
     n_startup_trials: int = 12,
     gamma: float = 0.25,
-    refine: bool = True,
     max_refine_iters: int = 5,
     log_path: str | None = None,
     path_anchor_axes: tuple[str, ...] = (),
     finite_check_args: tuple[str, ...] = (),
     cache_results: bool = True,
     reset_to_zero=None,
@@ -905,19 +940,27 @@ def bayesian_autotune(
     """Decorator mirroring ``@triton.autotune``. Extra kwargs:
     n_trials:                 successfully measured configs per key (TPE budget; configs
                               that fail to compile/run are skipped without consuming it)
     n_startup_trials:         random seed configs before the TPE model kicks in
     gamma:                    top fraction of measured configs treated as "good"
-    refine, max_refine_iters: coordinate-descent refinement after the TPE
     path_anchor_axes:         branch axes (kwarg names) whose values relocate the tile
                               optimum — one guaranteed max-tile anchor per value combo
                               (see _basin_anchor_indices); the kernel declares them, the
                               tuner stays independent of configuration details
     max_failures:             abort the tune once this many configs fail to compile/run before
-                              the measured budget is met (default n_trials; env override
-                              FINEGRAINED_AUTOTUNE_MAX_FAILURES). Without it a broadly broken
                               grid spins: failures don't consume the trial budget, so the search
                               keeps picking and each pick costs an O(configs) rescan.
-    log_path:                 JSONL of benched configs (or FINEGRAINED_AUTOTUNE_LOG)
     cache_results:            persist the tuned best config to disk (on by default) so later
                               runs skip the search+compile — see BayesianAutotuner.check_disk_cache"""
@@ -933,10 +976,12 @@ def bayesian_autotune(
             max_failures=max_failures,
             n_startup_trials=n_startup_trials,
             gamma=gamma,
-            refine=refine,
             max_refine_iters=max_refine_iters,
             log_path=log_path,
             path_anchor_axes=path_anchor_axes,
             finite_check_args=finite_check_args,
             cache_results=cache_results,
             **kwargs,

 import time
 from collections import defaultdict
 from contextlib import contextmanager
+from pathlib import Path
 import torch
     triton_key,
 )
+from .autotune_config import load_autotune_config
 logger = logging.getLogger(__name__)
+# Include the tuner implementation alongside Triton's kernel-source hash.
+_TUNER_SOURCE_HASH = hashlib.sha256(Path(__file__).read_bytes()).hexdigest()
+# Host-execution counter (``call_count_path`` in FINEGRAINED_AUTOTUNE_CONFIG): every
+# BayesianAutotuner.run() call increments its kernel's tally, dumped to that path at exit.
 # Discriminates cudagraph REPLAY (host code runs only at capture) from per-token
 # re-execution — the decode-gap forensics tool.
 _CALL_COUNTS = None
+if load_autotune_config().get("call_count_path"):
     import atexit
     _CALL_COUNTS = {}
+    def _dump_call_counts(path=load_autotune_config()["call_count_path"]):
         with open(path, "w") as f:
             json.dump(_CALL_COUNTS, f, indent=1, sort_keys=True)
         max_failures: int | None = None,
         n_startup_trials: int = 12,
         gamma: float = 0.25,
+        runoff: bool = True,
+        basin_probes: bool = True,
         max_refine_iters: int = 5,
         log_path: str | None = None,
         path_anchor_axes: tuple[str, ...] = (),
+        coupled_axes: tuple[tuple[str, ...], ...] = (),
         finite_check_args: tuple[str, ...] = (),
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
         self.fn_name = getattr(self.fn, "__name__", str(self.fn))
+        overrides = load_autotune_config()
+        # Bayesian trial budget — the per-decorator default, overridable via ``n_trials`` in
+        # FINEGRAINED_AUTOTUNE_CONFIG (quick sweeps / exhaustive runs without touching the
+        # decorators; set it >= grid size to fall back to stock exhaustive bench-all).
+        self.n_trials = overrides.get("n_trials", n_trials)
         # A tune aborts once this many configs have failed to compile/run without the measured
         # budget being met. Defaults to n_trials: a grid that cannot land its measurements without
         # that many rejects has a pruner gap, and the fix is to fence the dead region rather than
         # to tolerate the compiles. Its own axis so a kernel can raise it deliberately.
+        # an explicit argument wins over the config override, which wins over the default
+        if max_failures is None:  # ``strict`` aborts on the first failure
+            max_failures = 0 if overrides.get("strict") else self.n_trials
         self.max_failures = int(max_failures)
         self.n_startup_trials = n_startup_trials
         # top fraction of measured configs the TPE treats as "good"
         self.gamma = gamma
+        self.runoff = overrides.get("runoff", runoff)
+        self.basin_probes = overrides.get("basin_probes", basin_probes)
+        self.max_refine_iters = overrides.get("max_refine_iters", max_refine_iters)
         # Branch axes (by kwarg name) that relocate the TILE optimum — the caller's
         # declaration, since it is configuration knowledge (the tuner never hardcodes
         # axis names). Each distinct value combination of these axes gets one guaranteed
         # max-tile anchor; empty = median anchors only.
         self.path_anchor_axes = tuple(path_anchor_axes)
+        # Joint choices share one TPE density and one coordinate-refinement step.
+        self.coupled_axes = tuple(tuple(group) for group in coupled_axes)
+        coupled_names = [name for group in self.coupled_axes for name in group]
+        if any(len(group) < 2 for group in self.coupled_axes) or len(set(coupled_names)) != len(coupled_names):
+            raise ValueError("coupled_axes must contain disjoint groups of at least two configuration fields")
+        self._coupled_names = frozenset(coupled_names)
         # Output arg NAMES (a declaration — the tuner never hardcodes kernel details) whose
         # values are checked for NaN/inf after each config's bench: a config that writes
         # non-finite output scores inf and can never be crowned. The tuner times, it never
         # validates — this veto closes the NaN half of the "wrong-answer configs win
         # tunes" class at negligible cost (one isfinite pass per benched config).
         self.finite_check_args = tuple(finite_check_args)
+        # JSONL log of every benched (key, config, ms) — set here or via ``log_path`` in
+        # FINEGRAINED_AUTOTUNE_CONFIG. Analyse offline to prune bad configs.
+        self.log_path = log_path or overrides.get("log_path")
+        self.key_log_path = overrides.get("key_log_path")
+        self.tracebacks = overrides.get("tracebacks", False)
         # steady-state launch support: key extraction without dict builds (eager decode
         # is host-bound; the cache-hit path of run() is on the per-token critical path)
         self._arg_position = {name: i for i, name in enumerate(self.arg_names)}
             if root is not e:
                 msg += f" || root: {type(root).__name__}: {str(root)[-300:]}"
             err = f"{type(e).__name__}: {msg}"
+            if self.tracebacks:
                 logger.warning(
                     "[autotune] %s failed on %s:\n%s",
                     self.fn_name,
         self.nargs = None
         return ret
+    def _search_signatures(self, configs: List[Config]):
+        signatures = []
+        for config in configs:
+            values = config.all_kwargs()
+            signature = {(name,): value for name, value in values.items() if name not in self._coupled_names}
+            for group in self.coupled_axes:
+                try:
+                    signature[group] = tuple(values[name] for name in group)
+                except KeyError as exc:
+                    raise ValueError(f"Unknown coupled_axes configuration field: {exc.args[0]}") from exc
+            signatures.append(tuple(sorted(signature.items())))
+        return signatures
     def _bayesian_search(self, configs: List[Config], args, kwargs, key) -> Config:
         timings: Dict[int, float] = {}
+        sigs = self._search_signatures(configs)
         def bench_idx(idx: int) -> float:
             if idx in timings:
         # per-dimension model can write off a whole axis it never saw succeed (two dsv4 tunes
         # shipped 25-60% slow winners because their random seeds only sampled a basin's dead
         # configs), and coordinate descent can't recover a winner two coupled flips away.
+        n_startup = min(self.n_trials, max(2, self.n_startup_trials))
+        anchors = self._basin_anchor_indices(configs) if self.basin_probes else []
         order = list(range(len(configs)))
         random.Random(0).shuffle(order)
         warm_idx = self._warm_start_index(configs, key)
         # Coordinate-descent refinement: try single-dim perturbations around
         # the current best until no neighbor improves.
+        for _ in range(self.max_refine_iters):
+            best_idx = min(timings, key=timings.get)
+            best_sig = sigs[best_idx]
+            best_ms = timings[best_idx]
+            improved = False
+            for i, s in enumerate(sigs):
+                if i in timings:
+                    continue
+                diff = sum(1 for (_, a), (_, b) in zip(best_sig, s) if a != b)
+                if diff != 1:
+                    continue
+                if bench_idx(i) < best_ms:
+                    improved = True
+            if not improved:
+                break
         # Crown RUNOFF: the search's single-shot timings carry enough noise to flip crowns
         # between same-code tune rolls (gpt-oss prefill grouped rolled 1727-2150us on identical
         # the top candidates back-to-back — same warmed process, adjacent in time, so they see
         # the same clocks — and crown the winner of THAT comparison. Costs ~3 extra benches.
         finite = sorted((t, i) for i, t in timings.items() if t != float("inf"))
+        finalists = list(timings)
+        if self.runoff and len(finite) > 1:
+            finalists = [i for _, i in finite[:3]]
+            for i in finalists:
                 timings.pop(i, None)
                 bench_idx(i)  # memoizing closure: re-benches and re-records the popped entry
         self.configs_timings = {configs[i]: t for i, t in timings.items()}
+        # Other candidates still carry search-time samples, not comparable runoff measurements.
+        best = min(finalists, key=timings.get)
         if timings[best] == float("inf"):
             # crowning an arbitrary broken config would persist it as `best` on disk and
             # every later healthy process would replay it without re-tuning — the exact
             # sticky-contaminated-context scenario inf-forgiveness exists to survive
             self._report_bench_failures()
             raise RuntimeError(
+                f"[bayesian-autotune] {self.fn_name}: every final candidate failed "
+                "— nothing correct to crown or persist; the context may "
                 "be trap-contaminated (retry in a fresh process) or the grid is broken "
                 "for this launch (see the failure report above)."
             )
             pass
     def _basin_anchor_indices(self, configs: List[Config]) -> List[int]:
+        """Interior representatives of categorical basins, plus a wide output tile
+        per declared compute path. All anchors come from the pruned grid."""
         # Basin axes are DERIVED, not declared: a config kwarg with string or boolean
         # values is a branch axis (different code path — compute mode, operand swap,
         # warp specialization), which partitions the grid into disjoint performance
         if len(groups) <= 1:
             return []
         anchors = []
         coarse_seen = set()
         coarse_axes = [k for k in self.path_anchor_axes if k in basin_axes]
         for basin, idxs in groups.items():
+            # Choose the middle of every ordinal axis together. Lexicographic sorting
+            # ignores M and can put the "median" at extreme K/warps/stages. Each axis is
+            # scored by its value's rank normalized to [0, 1], so axes with many values
+            # do not outweigh axes with few; the anchor is the config closest to 0.5 on
+            # all of them (smallest squared distance), i.e. the basin's interior.
+            values = [configs[i].all_kwargs() for i in idxs]
+            axes = {
+                k: sorted({v[k] for v in values})
+                for k, value in values[0].items()
+                if isinstance(value, (int, float)) and not isinstance(value, bool)
+            }
+            def distance(i):
+                args = configs[i].all_kwargs()
+                return sum(
+                    (options.index(args[k]) / max(1, len(options) - 1) - 0.5) ** 2
+                    for k, options in axes.items()
+                )
+            anchors.append(min(idxs, key=distance))
             if not self.path_anchor_axes:
+                continue
             coarse = tuple(basin[basin_axes.index(k)] for k in coarse_axes)
+            if coarse not in coarse_seen:
                 coarse_seen.add(coarse)
+                # The existing wide-tile anchor protects coupled output-tile ridges,
+                # with interior contraction/warps/stages rather than their largest corner.
+                def wide_order(i):
+                    # largest output tile first; interior on the other axes breaks ties
+                    kwargs = configs[i].kwargs
+                    return -kwargs.get("BLOCK_SIZE_M", 1) * kwargs.get("BLOCK_SIZE_N", 1), distance(i)
+                wide = min(idxs, key=wide_order)
+                if wide != anchors[-1]:
+                    anchors.append(wide)
         return anchors
     def _warm_start_index(self, configs: List[Config], key=None):
         """Persist the tuned best config to Triton's on-disk cache so a later run (or process)
         skips the whole minutes-long search+compile — most of which is Triton codegen, not
         benching. Keyed exactly like Triton's own autotune cache (triton version + backend +
+        kernel/tuner source hashes + invalidating env + key + the full config grid), so any
+        of those changing re-tunes.
         Unlike Triton's stock version we do NOT bail when configs carry a ``pre_hook`` (ours
         always do): we persist each config's ``all_kwargs()`` — serializable and pre_hook-free —
             triton_key(),
             make_backend(driver.active.get_current_target()).hash(),
             fn.cache_key,
+            _TUNER_SOURCE_HASH,
             str(sorted(env_vars.items())),
             # Key NAMES and the search budget are determinants too: tuning_key holds
             # values only (renaming a key or swapping two same-typed keys would collide),
+            # and the n_trials override is not in Triton's invalidating-env list —
             # without n_trials a 3-trial exploratory winner gets replayed at 100 trials.
             str(self.keys),
+            f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma},"
+            f"runoff={self.runoff},basin_probes={self.basin_probes},max_refine_iters={self.max_refine_iters},"
+            f"coupled_axes={self.coupled_axes}",
             str(tuning_key),
         ] + [str(c) for c in configs]  # str(Config) is pre_hook-free and process-stable
         cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest()
         cache = get_cache_manager(cache_key)
         file_name = f"{fn.__name__[:150]}.bayes_autotune.json"
+        keylog = self.key_log_path
         if keylog:
             with open(keylog, "a") as f:
                 f.write(json.dumps({
     Held while a key is tuned and its crown published, so concurrent TP/EP ranks converge on the
     same config instead of each crowning its own (see ``check_disk_cache``). Advisory ``flock`` on
     a sidecar file next to the crown; entirely best-effort — no cache dir, no ``fcntl``, or a lock
+    that cannot be taken within ``lock_timeout`` seconds (FINEGRAINED_AUTOTUNE_CONFIG; default 30 min, since the holder
     is doing a real tune) all fall through to tuning locally. Losing the lock costs redundant work,
     never correctness."""
     lock_file = None
         cache_dir = getattr(cache, "cache_dir", None)
         if cache_dir:
             os.makedirs(cache_dir, exist_ok=True)
+            deadline = time.time() + load_autotune_config().get("lock_timeout", 1800.0)
             lock_file = open(os.path.join(cache_dir, f"{file_name}.lock"), "w")
             while True:
                 try:
     max_failures: int | None = None,
     n_startup_trials: int = 12,
     gamma: float = 0.25,
+    runoff: bool = True,
+    basin_probes: bool = True,
     max_refine_iters: int = 5,
     log_path: str | None = None,
     path_anchor_axes: tuple[str, ...] = (),
+    coupled_axes: tuple[tuple[str, ...], ...] = (),
     finite_check_args: tuple[str, ...] = (),
     cache_results: bool = True,
     reset_to_zero=None,
     """Decorator mirroring ``@triton.autotune``. Extra kwargs:
     n_trials:                 successfully measured configs per key (TPE budget; configs
                               that fail to compile/run are skipped without consuming it)
+                              (runtime override: ``n_trials`` in FINEGRAINED_AUTOTUNE_CONFIG)
     n_startup_trials:         random seed configs before the TPE model kicks in
     gamma:                    top fraction of measured configs treated as "good"
+    runoff:                   remeasure the top candidates before choosing the winner
+                              (runtime override: ``runoff``)
+    basin_probes:             measure categorical basin anchors before the TPE
+                              (runtime override: ``basin_probes``)
+    max_refine_iters:         coordinate-descent iterations after the TPE (0 disables)
+                              (runtime override: ``max_refine_iters``)
     path_anchor_axes:         branch axes (kwarg names) whose values relocate the tile
                               optimum — one guaranteed max-tile anchor per value combo
                               (see _basin_anchor_indices); the kernel declares them, the
                               tuner stays independent of configuration details
+    coupled_axes:             groups of config fields treated as joint choices by TPE and
+                              refinement, including Config fields such as num_stages
     max_failures:             abort the tune once this many configs fail to compile/run before
+                              the measured budget is met (default n_trials; runtime override
+                              ``strict`` in FINEGRAINED_AUTOTUNE_CONFIG sets 0). Without it a broadly broken
                               grid spins: failures don't consume the trial budget, so the search
                               keeps picking and each pick costs an O(configs) rescan.
+    log_path:                 JSONL of benched configs (or ``log_path`` in FINEGRAINED_AUTOTUNE_CONFIG)
     cache_results:            persist the tuned best config to disk (on by default) so later
                               runs skip the search+compile — see BayesianAutotuner.check_disk_cache"""
             max_failures=max_failures,
             n_startup_trials=n_startup_trials,
             gamma=gamma,
+            runoff=runoff,
+            basin_probes=basin_probes,
             max_refine_iters=max_refine_iters,
             log_path=log_path,
             path_anchor_axes=path_anchor_axes,
+            coupled_axes=coupled_axes,
             finite_check_args=finite_check_args,
             cache_results=cache_results,
             **kwargs,

build/torch-cuda/compat.py CHANGED Viewed

@@ -502,6 +502,18 @@ def is_sm10x() -> bool:
     return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10
 @functools.cache
 def is_sm90() -> bool:
     """Whether the CURRENT device is Hopper (sm_90) — the target the measured NVIDIA-backend

     return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10
+@functools.cache
+def is_blackwell() -> bool:
+    """Whether the CURRENT device is any Blackwell family: sm_10x (datacenter), sm_11x
+    (Jetson Thor) or sm_12x (consumer). Wider than ``is_sm10x``: these share the FP4
+    conversion instructions but not TMEM scaled-MMA."""
+    return (
+        torch.cuda.is_available()
+        and torch.version.hip is None
+        and torch.cuda.get_device_capability()[0] in (10, 11, 12)
+    )
 @functools.cache
 def is_sm90() -> bool:
     """Whether the CURRENT device is Hopper (sm_90) — the target the measured NVIDIA-backend

build/torch-cuda/embedding.py ADDED Viewed

	@@ -0,0 +1,306 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Selected-row lookup for floating-point and integer quantized embedding tables."""
+import torch
+import triton
+import triton.language as tl
+from .compat import (
+    FP8_DTYPE,
+    MX_SCALE_GROUP_K,
+    NVFP4_SCALE_GROUP_K,
+    add_op_namespace_prefix,
+    compile_time_only_triton_op,
+    compile_time_only_triton_wrap,
+    device_context,
+)
+from .formats import ue8m0_as_uint8
+from .loading.scales import decode_group_scale
+from .quant import e2m1_cols_to_f32
+@triton.jit
+def _decode_weight(value, row, FORMAT: tl.constexpr):
+    """Stored weight elements to fp32, before scaling. FP4 unpacks two columns per byte;
+    INT4 takes the looked-up row's nibble (even rows low, odd rows high)."""
+    if FORMAT == "nvfp4" or FORMAT == "mxfp4":
+        value = e2m1_cols_to_f32(value.to(tl.uint8))
+    elif FORMAT == "int4":
+        nibble = (value.to(tl.uint8).to(tl.int32) >> ((row % 2) * 4)) & 15
+        value = ((nibble ^ 8) - 8).to(tl.float32)
+    else:
+        value = value.to(tl.float32)
+    return value
+@triton.jit
+def _load_scale(
+    S,
+    program,
+    row,
+    col,
+    valid,
+    row_valid,
+    D: tl.constexpr,
+    S_ROW: tl.constexpr,
+    S_COL: tl.constexpr,
+    SCALE_BLOCK_ROW: tl.constexpr,
+    SCALE_BLOCK_COL: tl.constexpr,
+    FORMAT: tl.constexpr,
+    BLOCK: tl.constexpr,
+    PACKING: tl.constexpr,
+):
+    """Decoded fp32 scale for each of the tile's ``BLOCK`` output values."""
+    if (
+        D % BLOCK == 0
+        and S_COL != 0
+        and BLOCK % SCALE_BLOCK_COL == 0
+        and FORMAT != "mxfp8"
+    ):
+        # Load one scale per group; MXFP8 favors its existing flat byte loads.
+        groups = (program % (D // BLOCK)) * (BLOCK // SCALE_BLOCK_COL)
+        groups += tl.arange(0, BLOCK // SCALE_BLOCK_COL)
+        scale = tl.load(
+            S + (row // SCALE_BLOCK_ROW) * S_ROW + groups * S_COL,
+            row_valid,
+            other=0.0,
+        )
+        scale = decode_group_scale(scale).to(tl.float32)
+        scale = tl.broadcast_to(
+            scale[:, None], (BLOCK // SCALE_BLOCK_COL, SCALE_BLOCK_COL)
+        ).reshape((BLOCK,))
+    else:
+        scale_offsets = (row // SCALE_BLOCK_ROW) * S_ROW
+        if S_COL != 0:
+            scale_offsets += (col * PACKING // SCALE_BLOCK_COL) * S_COL
+        scale = tl.load(
+            S + scale_offsets,
+            row_valid if S_COL == 0 else valid,
+            other=0.0,
+        )
+        scale = decode_group_scale(scale).to(tl.float32)
+        if PACKING == 2:
+            scale = tl.broadcast_to(scale, (BLOCK // PACKING,))
+            scale = tl.interleave(scale, scale)
+    return scale
+@triton.jit
+def _embedding_kernel(
+    W,
+    Ids,
+    S,
+    G,
+    Out,
+    N: tl.constexpr,
+    D: tl.constexpr,
+    V: tl.constexpr,
+    W_ROW: tl.constexpr,
+    W_COL: tl.constexpr,
+    S_ROW: tl.constexpr,
+    S_COL: tl.constexpr,
+    SCALE_BLOCK_ROW: tl.constexpr,
+    SCALE_BLOCK_COL: tl.constexpr,
+    FORMAT: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    # FP4 packs adjacent columns; INT4 packs adjacent vocabulary rows.
+    packing: tl.constexpr = 2 if (FORMAT == "nvfp4" or FORMAT == "mxfp4") else 1
+    program = tl.program_id(0).to(tl.int64)
+    offsets = program * (BLOCK // packing) + tl.arange(0, BLOCK // packing)
+    valid = offsets < N * (D // packing)
+    # Tiles contained in one row share its index and per-row scale across all lanes.
+    if D % BLOCK == 0:
+        indices = program // (D // BLOCK)
+    else:
+        indices = offsets // (D // packing)
+    valid_rows = indices < N
+    row = tl.load(Ids + indices, valid_rows, other=0).to(tl.int64)
+    in_bounds = (row >= 0) & (row < V)
+    tl.device_assert(in_bounds | ~valid_rows, "embedding index out of range")
+    valid = valid & in_bounds
+    col = offsets % (D // packing)
+    weight_row = row // 2 if FORMAT == "int4" else row
+    value = tl.load(W + weight_row * W_ROW + col * W_COL, valid, other=0.0)
+    value = _decode_weight(value, row, FORMAT)
+    scale = _load_scale(
+        S, program, row, col, valid, valid_rows & in_bounds, D, S_ROW, S_COL,
+        SCALE_BLOCK_ROW, SCALE_BLOCK_COL, FORMAT, BLOCK, packing,
+    )
+    value = value * scale
+    if FORMAT == "nvfp4":
+        value = value * tl.load(G).to(tl.float32)
+    output_offsets = program * BLOCK + tl.arange(0, BLOCK)
+    tl.store(Out + output_offsets, value, output_offsets < N * D)
+def _check_inputs(weight, indices, scale, global_scale, output_dtype, weight_format, block_size):
+    """Validate an ``embedding_lookup`` call and resolve its format, its scale (UE8M0 viewed
+    as bytes), the ``(width, vocab, scale strides, scale blocks)`` kernel layout and the
+    output dtype. See ``embedding_lookup`` for the accepted layouts."""
+    if weight_format is None:
+        if global_scale is not None:
+            weight_format = "nvfp4"
+        else:
+            weight_format = "int8" if weight.dtype == torch.int8 else "fp8"
+    if weight_format not in ("fp8", "nvfp4", "mxfp8", "mxfp4", "int8", "int4"):
+        raise ValueError(f"Unsupported embedding weight format: {weight_format}.")
+    if weight.ndim != 2 or indices.dtype not in (torch.int32, torch.int64):
+        raise ValueError("Expected a 2D embedding table and int32/int64 indices.")
+    if weight.device != indices.device or weight.device != scale.device:
+        raise ValueError("Embedding table, indices, and scales must share a device.")
+    packed_columns = weight_format in ("nvfp4", "mxfp4")
+    width = weight.shape[1] * (2 if packed_columns else 1)
+    vocab = weight.shape[0] * (2 if weight_format == "int4" else 1)
+    if width == 0:
+        raise ValueError("Embedding width must be positive.")
+    float_dtypes = (torch.float16, torch.bfloat16, torch.float32)
+    scale_row = scale_col = 0
+    row_block = col_block = 1
+    if weight_format in ("fp8", "int8"):
+        expected_dtype = FP8_DTYPE if weight_format == "fp8" else torch.int8
+        if weight.dtype != expected_dtype or scale.dtype not in float_dtypes:
+            raise ValueError(
+                f"{weight_format.upper()} requires {expected_dtype} weights and floating scales."
+            )
+        if scale.numel() == 1 and weight_format == "fp8":
+            pass
+        elif scale.shape in ((vocab,), (vocab, 1)):
+            scale_row = scale.stride(0)
+        elif (
+            weight_format == "fp8"
+            and block_size > 0
+            and scale.shape
+            == (triton.cdiv(vocab, block_size), triton.cdiv(width, block_size))
+        ):
+            row_block = col_block = block_size
+            scale_row, scale_col = scale.stride()
+        else:
+            raise ValueError(
+                "FP8 scales must be scalar, per-row, or a block grid; INT8 scales must be per-row."
+            )
+    else:
+        if weight_format == "mxfp8":
+            if weight.dtype != FP8_DTYPE:
+                raise ValueError("MXFP8 requires E4M3 weights.")
+        elif weight.dtype not in (torch.int8, torch.uint8):
+            raise ValueError("FP4 and INT4 weights must be packed int8/uint8 bytes.")
+        if weight_format == "nvfp4":
+            col_block = NVFP4_SCALE_GROUP_K
+            scale_dtypes = (FP8_DTYPE,)
+        elif weight_format in ("mxfp8", "mxfp4"):
+            col_block = MX_SCALE_GROUP_K
+            scale = ue8m0_as_uint8(scale)
+            scale_dtypes = (torch.uint8,)
+        else:
+            if scale.ndim != 2 or scale.shape[1] == 0 or width % scale.shape[1]:
+                raise ValueError(
+                    "INT4 scales must be a per-row grid of equal column groups."
+                )
+            col_block = width // scale.shape[1]
+            scale_dtypes = float_dtypes
+        if (
+            width % col_block
+            or scale.shape != (vocab, width // col_block)
+            or scale.dtype not in scale_dtypes
+        ):
+            raise ValueError(
+                f"{weight_format.upper()} requires a ({vocab}, {width // col_block}) scale grid with dtype in {scale_dtypes}."
+            )
+        scale_row, scale_col = scale.stride()
+    if weight_format == "nvfp4":
+        if (
+            global_scale is None
+            or global_scale.dtype != torch.float32
+            or global_scale.numel() != 1
+            or global_scale.device != weight.device
+        ):
+            raise ValueError(
+                "NVFP4 requires one FP32 global multiplier on the table's device."
+            )
+    elif global_scale is not None:
+        raise ValueError("Only NVFP4 embeddings use a global multiplier.")
+    dtype = output_dtype or (
+        torch.bfloat16 if weight_format in ("nvfp4", "mxfp8", "mxfp4") else scale.dtype
+    )
+    if dtype not in float_dtypes:
+        raise ValueError("Embedding output must be float16, bfloat16, or float32.")
+    return weight_format, scale, (width, vocab, scale_row, scale_col, row_block, col_block), dtype
+@compile_time_only_triton_op(add_op_namespace_prefix("embedding_lookup"))
+def embedding_lookup(
+    weight: torch.Tensor,
+    indices: torch.Tensor,
+    scale: torch.Tensor,
+    global_scale: torch.Tensor | None = None,
+    output_dtype: torch.dtype | None = None,
+    *,
+    weight_format: str | None = None,
+    block_size: int = 128,
+) -> torch.Tensor:
+    """Gather and dequantize selected rows without expanding the embedding table.
+    ``weight_format`` selects ``fp8``, ``nvfp4``, ``mxfp8``, ``mxfp4``, ``int8``, or ``int4``.
+    By default, a global scale selects NVFP4, signed int8 weights select INT8, and otherwise
+    FP8 is used. FP8 accepts scalar, per-row, or square block scales (``block_size``).
+    NVFP4 uses E4M3 group-16 scales and one FP32 multiplier; MX formats use UE8M0 group-32
+    scales. FP4 packs adjacent columns, low nibble first. INT8 uses per-row scales; INT4
+    packs adjacent vocabulary rows and uses per-row column groups inferred from its scale
+    grid. Integer values are signed, symmetric, and have no zero point.
+    Outputs default to the floating scale dtype, or bf16 for NVFP4/MX. Table/grid strides
+    may describe TP slices. Indices follow ``torch.nn.functional.embedding``: negative and
+    out-of-range values fail; padding indices look up their stored row. Empty, repeated,
+    and noncontiguous indices work.
+    """
+    weight_format, scale, layout, dtype = _check_inputs(
+        weight, indices, scale, global_scale, output_dtype, weight_format, block_size
+    )
+    width, vocab, scale_row, scale_col, row_block, col_block = layout
+    out = torch.empty((*indices.shape, width), dtype=dtype, device=weight.device)
+    rows, numel = indices.numel(), out.numel()
+    if rows == 0:
+        return out
+    # Keep small lookups parallel; larger tiles amortize scheduling for bulk lookups.
+    block = 256
+    if numel >= 256 * 1024:
+        block = 1024
+    if numel >= 4 * 1024 * 1024:
+        block = 4096
+    with device_context(weight.device):
+        compile_time_only_triton_wrap(_embedding_kernel)[(triton.cdiv(numel, block),)](
+            weight,
+            indices.reshape(-1),
+            scale,
+            global_scale,
+            out,
+            rows,
+            width,
+            vocab,
+            weight.stride(0),
+            weight.stride(1),
+            scale_row,
+            scale_col,
+            row_block,
+            col_block,
+            weight_format,
+            block,
+            num_warps=4,
+            enable_fp_fusion=False,
+            debug=True,
+        )
+    return out

build/torch-cuda/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "finegrained-kernels",
-  "id": "_finegrained_kernels_cuda_43f1581",
   "version": 0,
   "kernels-minver": "0.14.0",
   "license": "Apache-2.0",
@@ -12,13 +12,15 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "43nBPHxybam5ChWZmsz1Ji+pjJ6K6zF5lzwbqPcMwE4=",
-      "_ops.py": "OSOJZM6uXgutp0kDIGNXhiJDHbTZg6slWOmcmwDfRB4=",
       "backward.py": "K1S6KD3J/12SthXiqBcmNO9kUikctQpccsbxsO8lpRg=",
       "batched.py": "U4zfmMlnQjr4j+kiIxzzohKmIXG0rtQX38c3i+aBgT8=",
-      "bayesian_autotuner.py": "63cy/g/2qR2hwuRYgIoM05Ylhq7QCncazpaBTHd4iGs=",
-      "compat.py": "4z8Q3urWNhfDM7qeg9gSZslVk9bQl5e46Mc2lT8SwpQ=",
       "descriptors.py": "4acRMazGPFlbydNt1Hbo8RpAqb3fOtnOQBv5yb4onqw=",
       "epilogue.py": "vIeMt0EqWgpvadDw5r38/5TVRMMzIMthfXGHPa3lTeY=",
       "formats.py": "rtzJgKnIK1PV/sqNnPGO67Ve6czzx1ldyIhs9HhydSk=",
       "grouped.py": "jv+qffb9l+MRDbo9UdA1yV15o8QcgNvE6ZYzCpPrpYM=",
@@ -29,11 +31,11 @@
       "mma.py": "gLkCU5QjhQVMacixNR24Yi07aNelQyN+I0pJOuRxE4E=",
       "moe.py": "4n/TzWy0kRXAL9qGY/aChbdSdCpahVVmMUKCfd6/zF0=",
       "norm.py": "VvE6GUE0Oztf7PDcviLCcB4fhnXt7/zgdE+aApv6DBM=",
-      "pruners.py": "CzVRCpThHvmt4N3YH0GJ+JV/Cx6c+bV6vJm1lXyCdtM=",
-      "quant.py": "Vqon6Gl/VsWZ+5ej1FDhD9aqFQUBgUb0HKXqrk5CoFQ=",
       "scheduling.py": "cxwoL5NTJxARVxLStDsNalm/jqbkCObm8uyLRFL/yYc=",
       "swizzle.py": "ueIJ9FsO8QZhVt3nMFlA+8b+U9/ZHabHFlhufsONV04=",
-      "symbols.json": "pajgmx0vGKA1tt5P+W3f88iV3+m3SWIYcnFhZ+YrLBI="
     }
   },
   "provenance": {
@@ -43,7 +45,7 @@
       "dirty": false
     },
     "kernel": {
-      "commit": "43f15810b3acaad1d17f67745f0676896b6fcf0d",
       "dirty": false
     }
   }

 {
   "name": "finegrained-kernels",
+  "id": "_finegrained_kernels_cuda_ec5cc90",
   "version": 0,
   "kernels-minver": "0.14.0",
   "license": "Apache-2.0",
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "xFRwaLTont8ZLTxS5ajuiDqM9gWwXAU4xap2prqdhrY=",
+      "_ops.py": "DwdUBjZmrNFv1HEdJDNcSfPqBc4da6D3M8v2dUzUA3g=",
+      "autotune_config.py": "d5fb2Vqss0fgC9TZbarrmQIInL1n1L7rHz6o3Gqv6cI=",
       "backward.py": "K1S6KD3J/12SthXiqBcmNO9kUikctQpccsbxsO8lpRg=",
       "batched.py": "U4zfmMlnQjr4j+kiIxzzohKmIXG0rtQX38c3i+aBgT8=",
+      "bayesian_autotuner.py": "NDRUQrhk+bEEvrOkIilWJPfqdKpcsd4E4QQ9AHzylKI=",
+      "compat.py": "K3RCq1ETi80duzv3zSldIwrL5e46vZAWtLxLlg1LXvM=",
       "descriptors.py": "4acRMazGPFlbydNt1Hbo8RpAqb3fOtnOQBv5yb4onqw=",
+      "embedding.py": "k7jiz2vtEJuYKo+w6TuS/4WO1D4Bv96O5DMFjEpYUGo=",
       "epilogue.py": "vIeMt0EqWgpvadDw5r38/5TVRMMzIMthfXGHPa3lTeY=",
       "formats.py": "rtzJgKnIK1PV/sqNnPGO67Ve6czzx1ldyIhs9HhydSk=",
       "grouped.py": "jv+qffb9l+MRDbo9UdA1yV15o8QcgNvE6ZYzCpPrpYM=",
       "mma.py": "gLkCU5QjhQVMacixNR24Yi07aNelQyN+I0pJOuRxE4E=",
       "moe.py": "4n/TzWy0kRXAL9qGY/aChbdSdCpahVVmMUKCfd6/zF0=",
       "norm.py": "VvE6GUE0Oztf7PDcviLCcB4fhnXt7/zgdE+aApv6DBM=",
+      "pruners.py": "5kZlpQiAvhdy8tiD8+tsy3YLDwYpHGY/aKAsxIVkTRw=",
+      "quant.py": "1fNRvb63DAOcFNLfKw3cSwMqSwLua4smI+xvhlsUk3k=",
       "scheduling.py": "cxwoL5NTJxARVxLStDsNalm/jqbkCObm8uyLRFL/yYc=",
       "swizzle.py": "ueIJ9FsO8QZhVt3nMFlA+8b+U9/ZHabHFlhufsONV04=",
+      "symbols.json": "sVGkXI22bMHsNCUTsKl1Hj70D7j8D+/KhJ2CS1JhIPA="
     }
   },
   "provenance": {
       "dirty": false
     },
     "kernel": {
+      "commit": "ec5cc90e1657566c66c62b5021435edbcfa348c3",
       "dirty": false
     }
   }

build/torch-cuda/metadata.json.sigstore CHANGED Viewed

@@ -1 +1 @@

- {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHKjCCBrGgAwIBAgIUYjRfDt64EwMm2w5mvJf2snFI03gwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDIwMDc0NTA4WhcNMjYxMDIwMDc1NTA4WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEVosgukER4CXhkkdfSXj6fQjeiWig2uOUbciDoDxHWLHj/LYrEK37NIX0hVju279Sg0dyPpiIctPzCStUj3ecUaOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBCg0M2YxNTgxMGIzYWNhYWQxZDE3ZjY3NzQ1ZjA2NzY4OTZiNmZjZjBkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDCg0M2YxNTgxMGIzYWNhYWQxZDE3ZjY3NzQ1ZjA2NzY4OTZiNmZjZjBkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDCg0M2YxNTgxMGIzYWNhYWQxZDE3ZjY3NzQ1ZjA2NzY4OTZiNmZjZjBkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoNDNmMTU4MTBiM2FjYWFkMWQxN2Y2Nzc0NWYwNjc2ODk2YjZmY2YwZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTM4Nzk5OTk0Ny9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhWEVw1gAABAMARjBEAiBQIWsZSWB9SJnGRw0TAla0o5FB0zYAm05drx5ceF8hZwIgcJXjQg1k+AOW7IQ8jSl/xqcqBqMtOZZ4ZDUyAW74/3EwCgYIKoZIzj0EAwMDZwAwZAIwQUUPe2Wu9HMn7udhGmysCY3wmjvsq5XAq26BaF2dsHqb654f7eVjlzgv69NyN+INAjBS+0uAAoIQtkuKGmdxafnxXTlaVHKjkgGeHzCOM9vQ2d1M5Pe6SMCzIJX6Be0kZv8="},"tlogEntries":[{"logIndex":"3333276781","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1792482308","inclusionPromise":{"signedEntryTimestamp":"MEQCIBFaOIQqgtBjrkEZL6N9d+ec1YTSMPhobCkwlT0b0kSCAiA1RnJFqAyg2V4NirN3+K9sq3AlLM6Fr7y0yLF5WUaPJg=="},"inclusionProof":{"logIndex":"3211372519","rootHash":"AAuNOep3f6B6LEiLxEJLdf3MrE+3zqg7zm9NZq+xwfg=","treeSize":"3211372521","hashes":["kSi6E1FLwjGxB6aGLjryVwfihvJqSpRS8HNrf320SFM=","KpdKzAVftSQYDDBeFgeUVygXUujlCunLuEufgAMcTCs=","HsnVOmvwSYcTBx5ev7Cwx82EDIzB4cL99kQg3bhkmk8=","ZO4WKHe6e+kjduSIqG+P39aU3jpzqV5WS46QbUt/1lM=","DSCf6YnSfwKGZbSQHikozOUYjJVNknivnUg1PC0RRns=","OMrFmT5b3Baa5tQ18hoVjZoPhI3C49iB09L+Byp3z2o=","/2z6rGxIRvEQCFmmSVuv3acTBVHwi852tRBHlC4/vSE=","aHvYjdQ34DQYMV9K0YMPPiIHQk74oGCBXn6ekwMGA3U=","pPbxA70P/0HkEiYY2jZB61Bpv0868ThbYHS5oemCUxo=","eKD1kuBv/uuXm3Q+ddDIoOM07eGyxtScSe41MjUUwEA=","35bnJ5pQ9Q/ajuvM0i7P6+YcssTBPuuptjhAKxuEmfo=","OcDrZ0RIsie8PY2uv3Av6cECp65xfWyxKOKcM1/c3sc=","zj28xnvSQbA7lgK1R588ZHvZYy6Gy03wZvsJNVfxG9o=","2ljWFy3+P4yeeVU2ALgF4M45of0N5Rv4W1Z/jm/9n5w=","3aUq4N47Z1/AszcODIxRgusLXNRFZAnHFg6Vur7t22o=","PwdgHXfl6JMsozFa2gp9OjWOM8+7eYRrqPJa9FEI6k4=","HovTPBWzoErvUg1u5sIi3x1ZonZ7yJekVsBrmpomCIA=","G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=","qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=","gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=","NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=","qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=","xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3211372521\nAAuNOep3f6B6LEiLxEJLdf3MrE+3zqg7zm9NZq+xwfg=\n\n— rekor.sigstore.dev wNI9ajBGAiEAiHqnOXsqc9SUah3w9DEbIpUb2csqmTT7SpdakR2wOmYCIQDAikNdot3q0+e2IXLza+Diw4Tn08v5eti/aHYhGvKzTg==\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJlMTE4OGEyZmEzYWM5NGJkYTFmMzEzMDQ1NDU2YzFhYmNmZTcwNmM1ZDU4YmMwN2RjZjA5Yjg3NmZlZjM0N2UwIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNscUdyeDRIRTRNdG1Uc24rZFZWTVEyWUEweTRjUGpLT0xBY1E3emlENGNnSWdBOEt4YU5qYTV5RlZkcXoxRG5qM0tpbE8yZDg3WHdpWmU2NExMd3B3c2FFPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExha05EUW5KSFowRjNTVUpCWjBsVldXcFNaa1IwTmpSRmQwMXRNbmMxYlhaS1pqSnpia1pKTURObmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1hkTlJHTXdUbFJCTkZkb1kwNU5hbGw0VFVSSmQwMUVZekZPVkVFMFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZXYjNObmRXdEZValJEV0docmEyUm1VMWhxTm1aUmFtVnBWMmxuTW5WUFZXSmphVVFLYjBSNFNGZE1TR292VEZseVJVc3pOMDVKV0RCb1ZtcDFNamM1VTJjd1pIbFFjR2xKWTNSUWVrTlRkRlZxTTJWalZXRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmN3VFRKWmVBcE9WR2Q0VFVkSmVsbFhUbWhaVjFGNFdrUkZNMXBxV1ROT2VsRXhXbXBCTWs1NldUUlBWRnBwVG0xYWFscHFRbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEWnpCTk1sbDRUbFJuZUUxSFNYcFpWMDVvQ2xsWFVYaGFSRVV6V21wWk0wNTZVVEZhYWtFeVRucFpORTlVV21sT2JWcHFXbXBDYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEWnpCTk1sbDRDazVVWjNoTlIwbDZXVmRPYUZsWFVYaGFSRVV6V21wWk0wNTZVVEZhYWtFeVRucFpORTlVV21sT2JWcHFXbXBDYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlPUkU1dFRWUlZORTFVUW1sTk1rWnFXVmRHYTAxWFVYaE9NbGt5VG5wak1FNVhXWGRPYW1NeVQwUnJNbGxxV20wS1dUSlpkMXBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUk5ORTU2YXpWUFZHc3dUbms1YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVjBWV2R6Rm5RVUZDUVUxQlVtcENSVUZwUWxGSlYzTmFVMWRDT1ZOS2JrZFNkekJVUVd4aE1HODFSa0l3ZWxsQmJUQTFaQXB5ZURWalpVWTRhRnAzU1dkalNsaHFVV2N4YXl0QlQxYzNTVkU0YWxOc0wzaHhZM0ZDY1UxMFQxcGFORnBFVlhsQlZ6YzBMek5GZDBObldVbExiMXBKQ25wcU1FVkJkMDFFV25kQmQxcEJTWGRSVlZWUVpUSlhkVGxJVFc0M2RXUm9SMjE1YzBOWk0zZHRhblp6Y1RWWVFYRXlOa0poUmpKa2MwaHhZalkxTkdZS04yVldhbXg2WjNZMk9VNTVUaXRKVGtGcVFsTXJNSFZCUVc5SlVYUnJkVXRIYldSNFlXWnVlRmhVYkdGV1NFdHFhMmRIWlVoNlEwOU5PWFpSTW1ReFRRbzFVR1UyVTAxRGVrbEtXRFpDWlRCclduWTRQUW90TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgY3k0z/Mt7jMBEhOjvjf2G4CnoToAJ6VFGWkxAh1uoOYCFQCJXH8pBrNsoDup37wFNKA026PnoBgPMjAyNjEwMjAwNzQ1MDhaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MTAyMDA3NDUwOFowLwYJKoZIhvcNAQkEMSIEIF71Fa9Y2eBvEIVLLboLwnryQ7XILn3/JFn/+1wF+4ddMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMQDjXGEzJIPO0LhMOvGESo6oPJAF8cA/W/LS4FxKIXt2lA+oMCsX4UlFfHVmyUHEi8MCMGGJHGkEaw0eWLT2aTvt7BArmK0TWJXsTkbdnXPHHieoVQ6XUIfg+a76JBGvPIZN0w=="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"4RiKL6OslL2h8xMEVFbBq8/nBsXVi8B9zwm4dv7zR+A="},"signature":"MEUCIQClqGrx4HE4MtmTsn+dVVMQ2YA0y4cPjKOLAcQ7ziD4cgIgA8KxaNja5yFVdqz1Dnj3KilO2d87XwiZe64LLwpwsaE="}}

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHKzCCBrGgAwIBAgIUP/mmZM1RBNZy+2BaxIPboKaANfUwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDIwMDk1MzU0WhcNMjYxMDIwMTAwMzU0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEcYog7arHY0Y9LB1dbz4EMiisjAeogYrVaM/azAT620TxPmQwSkpj3gM5WuntjYLVKH7v6Y/ZJKbtEH3Iyz0eQqOCBdAwggXMMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChlYzVjYzkwZTE2NTc1NjZjNjZjNjJiNTAyMTQzNWVkYmNmYTM0OGMzMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChlYzVjYzkwZTE2NTc1NjZjNjZjNjJiNTAyMTQzNWVkYmNmYTM0OGMzMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChlYzVjYzkwZTE2NTc1NjZjNjZjNjJiNTAyMTQzNWVkYmNmYTM0OGMzMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZWM1Y2M5MGUxNjU3NTY2YzY2YzYyYjUwMjE0MzVlZGJjZmEzNDhjMzAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTQwMzQ1ODE5OS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGJBgorBgEEAdZ5AgQCBHsEeQB3AHUA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhWLtUgAAABAMARjBEAiAe/acaJEWfp83gcFkDKvoFLXrzTxjWugLIDNpSFeLusgIgI6Pe4FEqWllfViiFTAQVVYJoN4GbF+Zqar2IqNRKXywwCgYIKoZIzj0EAwMDaAAwZQIwTGwKwrp9v3bnF+jcmcxJDQlNH+G/X094ztAojPb9XMoglDopWFLG5CCHJwTVDBcxAjEAwAHPN6qnPDCH1MsFN1mUa25cXwEcCnBKYfOKJuuV4GuJ48l7lBLHt/XgHMSKAp0o"}, "tlogEntries":[{"logIndex":"3335086358", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792490034", "inclusionPromise":{"signedEntryTimestamp":"MEQCIDh5l4rIduFEIESVLdxB7yNjOmzLB6wHWBP+dLlyoALjAiBlffoVN9/Y6hMJ3jh8k9Tx78uilgTkTl+MFakp/tnDEw=="}, "inclusionProof":{"logIndex":"3213182096", "rootHash":"ETOUthJxn7Z36oDbegdN4ukdMMDtHTet6l/TBkV1Hww=", "treeSize":"3213182167", "hashes":["N4zOdeYV9bT+mL8n5/1CGmSuA2VV+UhN9MDbMneloJw=", "mj7NwwalWVWDX0Kh+to0lJbH1rGZDbNf4HcF9eItBVM=", "fwnSmRcMGxhsozTLY4LcxDpo08kcNy/vEPBGeeBSFbk=", "Cj68CvJffFHJ7/61leQDCepCbXGoBCS1YUvSGi3bm/Q=", "lfASF7qNkK6GWuggPpPvvjXCM94C8NJa0Wot/arlKXc=", "Yt7IuEE3A6hytfQiIcf8HBbHvrwdyXrEKDoszquugMM=", "Wtj/zv8pbDRBXP70LrlX25kih3b6H7UquUi/a+h1UG0=", "MgGbTQdHMx9a+WZnHvQoG7hXqexepfCSgGbCVHQdTME=", "5VvnPwgJg7mEz2CpXW/bEuCEagORQRlWh9Ndtihxwi0=", "PHxQAH8g/7NVsymUZ0Te3puxjBd3DEJNVas6rZCulgs=", "rYA6lUjt9opyITSlS0ZGyFtQo1JeGpc94KBmcxBHawE=", "Zz0ghh+5MBl4Vee4jXWbIrdb3i68SxmtMzhdMeM0hOg=", "6Rb7uOgk2v5YKcFSbAPOpLlWAIN1ciDPdlp4J/CUu5k=", "HovTPBWzoErvUg1u5sIi3x1ZonZ7yJekVsBrmpomCIA=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3213182167\nETOUthJxn7Z36oDbegdN4ukdMMDtHTet6l/TBkV1Hww=\n\n— rekor.sigstore.dev wNI9ajBGAiEA136ugN4DB8HrrSrQe1IkuD6gcEvjcYgnISn5rjXpld4CIQC11ZBPQ1jo/Uz+0/lcjMNx+Lkg4q0KC5vUVkoOyQojQA==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJhNDA1MWRmMTQ5MDNiM2Y4OTZmMzg3NWQyOTlmNTU4OTU5YjZmZWVjNzc5YmQwOWJjM2M5YWUxYjUxNmZhNjhjIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUQ4RVVFMEhWUHRzdGlVdDRFR0pnVG5ocitmUU9uRTg4SnpMYmVQN1M3MmV3SWdIZ0xaUy9lZFRDM1dMTmsrazdHQUtyUTJxeVdHZk1od3NVaXVQd2tEQVQ0PSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaExla05EUW5KSFowRjNTVUpCWjBsVlVDOXRiVnBOTVZKQ1RscDVLekpDWVhoSlVHSnZTMkZCVG1aVmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFU1hkTlJHc3hUWHBWTUZkb1kwNU5hbGw0VFVSSmQwMVVRWGROZWxVd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZqV1c5bk4yRnlTRmt3V1RsTVFqRmtZbm8wUlUxcGFYTnFRV1Z2WjFseVZtRk5MMkVLZWtGVU5qSXdWSGhRYlZGM1UydHdhak5uVFRWWGRXNTBhbGxNVmt0SU4zWTJXUzlhU2t0aWRFVklNMGw1ZWpCbFVYRlBRMEprUVhkbloxaE5UVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmhzV1hwV2FncFplbXQzV2xSRk1rNVVZekZPYWxwcVRtcGFhazVxU21sT1ZFRjVUVlJSZWs1WFZtdFpiVTV0V1ZSTk1FOUhUWHBOUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUd4WmVsWnFXWHByZDFwVVJUSk9WR014Q2s1cVdtcE9hbHBxVG1wS2FVNVVRWGxOVkZGNlRsZFdhMWx0VG0xWlZFMHdUMGROZWsxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUd4WmVsWnFDbGw2YTNkYVZFVXlUbFJqTVU1cVdtcE9hbHBxVG1wS2FVNVVRWGxOVkZGNlRsZFdhMWx0VG0xWlZFMHdUMGROZWsxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhVjAweFdUSk5OVTFIVlhoT2FsVXpUbFJaTWxsNldUSlplbGw1V1dwVmQwMXFSVEJOZWxac1drZEthbHB0UlhvS1RrUm9hazE2UVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUlJkMDE2VVRGUFJFVTFUMU01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkS0NrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NITkZaVkZDTTBGSVZVRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVjB4MFZXZEJRVUZDUVUxQlVtcENSVUZwUVdVdllXTmhTa1ZYWm5BNE0yZGpSbXRFUzNadlJreFljbnBVZUdwWGRXZE1TUXBFVG5CVFJtVk1kWE5uU1dkSk5sQmxORVpGY1Zkc2JHWldhV2xHVkVGUlZsWlpTbTlPTkVkaVJpdGFjV0Z5TWtseFRsSkxXSGwzZDBObldVbExiMXBKQ25wcU1FVkJkMDFFWVVGQmQxcFJTWGRVUjNkTGQzSndPWFl6WW01R0sycGpiV040U2tSUmJFNUlLMGN2V0RBNU5IcDBRVzlxVUdJNVdFMXZaMnhFYjNBS1YwWk1SelZEUTBoS2QxUldSRUpqZUVGcVJVRjNRVWhRVGpaeGJsQkVRMGd4VFhOR1RqRnRWV0V5TldOWWQwVmpRMjVDUzFsbVQwdEtkWFZXTkVkMVNnbzBPR3czYkVKTVNIUXZXR2RJVFZOTFFYQXdid290TFMwdExVVk9SQ0JEUlZKVVNVWkpRMEZVUlMwdExTMHRDZz09In19fX0="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgfFMO/p/NSekrYglMmlihkzviE2/TMjI7I9gvE+jPVZwCFHtSTryVjqRDdM+RmwtYaQ+zgsgnGA8yMDI2MTAyMDA5NTM1NFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDIwMDk1MzU0WjAvBgkqhkiG9w0BCQQxIgQg3f6Q/nF2OKTQsvL152ZutBpL80L7gdl0XBjdZCcouJ8wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxAKnEXDu2NvqZyacw0ZGY4AbtlHSBesVRge3qaPnSC2bWDBInFL34StY/aRcVjYTU0AIwZZOS9rSnZ5Y96uJwvpNIowj14zfaPu6vOZLCfkwkTorS0cyiyGuRF8lzjClRukKm"}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"pAUd8UkDs/iW84ddKZ9ViVm2/ux3m9Cbw8muG1Fvpow="}, "signature":"MEUCIQD8EUE0HVPtstiUt4EGJgTnhr+fQOnE88JzLbeP7S72ewIgHgLZS/edTC3WLNk+k7GAKrQ2qyWGfMhwsUiuPwkDAT4="}}

build/torch-cuda/pruners.py CHANGED Viewed

@@ -253,7 +253,7 @@ def gated_pointer_weight_warp_spec_pruner():
     sub-128 M tile — plain pointer/pointer at BM=128 and w4/w8 compiles and is often the winner.
     Region measured by tuning the GPT-OSS W4A16 problem (N=K=2880) with a fresh Triton cache and
-    ``FINEGRAINED_AUTOTUNE_LOG``: of 129 WS configs, the 12 that failed are exactly the ones this
     predicate names, and it drops none of the 117 that compiled. Every failure sat in the GATE
     shape — the ungated shape's 63 WS configs all compiled. Unfenced these are wasted compiles per
     new shape (the tuner memoizes them only after paying once). The ``BM < 128`` clause rests on a

     sub-128 M tile — plain pointer/pointer at BM=128 and w4/w8 compiles and is often the winner.
     Region measured by tuning the GPT-OSS W4A16 problem (N=K=2880) with a fresh Triton cache and
+    an autotune ``log_path``: of 129 WS configs, the 12 that failed are exactly the ones this
     predicate names, and it drops none of the 117 that compiled. Every failure sat in the GATE
     shape — the ungated shape's 63 WS configs all compiled. Unfenced these are wasted compiles per
     new shape (the tuner memoizes them only after paying once). The ``BM < 128`` clause rests on a

build/torch-cuda/quant.py CHANGED Viewed

@@ -19,7 +19,7 @@ from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait
 from .bayesian_autotuner import bayesian_autotune
 from .formats import global_scale_stride, is_per_expert_global
-from .compat import FP8_DTYPE, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, compile_time_only_triton_wrap, device_context, is_sm10x, pdl_launch_kwargs
 from .swizzle import swizzle_store_block
 from .scheduling import build_tile_layout, resolve_tile_inline
@@ -61,16 +61,14 @@ def fp8_act_quant_inline(
     return a_fp8, a_scale
-# cvt.e2m1x2.f32 (hardware FP4 pack) exists only on sm_100 (Blackwell). Resolved once at
-# import as a compile-time constexpr for the jit helper below; the ALU fallback compiles
-# everywhere else. ``is_sm10x`` is driverless-safe, so no import-time guard is needed.
-_E2M1_HW_CVT = tl.constexpr(is_sm10x())
 @triton.jit
 def _quant_e2m1_packed(v, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):
     """Pack signed, pre-scaled values on the E2M1 grid to ``(M, K//2)`` uint8 (first value
-    of each pair in the low nibble). On sm_100 the Blackwell hardware convert
     ``cvt.rn.satfinite.e2m1x2.f32`` (two f32 → one packed byte, first operand → HIGH nibble)
     does it in one instruction; elsewhere a ``>=``-threshold bucketize builds the code. The
     two agree except at exact E2M1 midpoints (0.25, 0.75, …), where the hardware rounds to
@@ -187,7 +185,7 @@ def _e2m1_code_to_bf16_bits(code):
 @triton.jit
 def _e2m1_pair_to_f16(packed):
     """One packed E2M1 byte -> the two fp16 values it holds, as a uint32 (low half = low nibble),
-    through the sm_100 CVT unit: ONE instruction per pair against the bit-twiddle form's ~10
     integer ops. Every E2M1 magnitude is exact in fp16, so this is a value-identical decode, and
     the direction of the unpack is irrelevant to the cost — the arithmetic is the whole bill
     (measured on the gpt-oss W4A16 decode weight stream: 39.6us bit-twiddle, either axis, vs
@@ -210,7 +208,7 @@ def e2m1_to_bf16(b_packed):
     interleave along K with a SINGLE 3D axis-swap (``join`` -> ``[R, C, 2]`` -> swap last two ->
     reshape) instead of the three 2D transposes ``e2m1_to_e4m3`` needs — the weight-only unpack hot path
     (fp4 weight upcast in-loop, the caller widening to the activation dtype). K order low-nibble-first, matching
-    ``e2m1_to_e4m3``: ``[byte0_lo, byte0_hi, byte1_lo, ...]``. On sm_100 the codes decode through
     the hardware CVT (``_e2m1_pair_to_f16``, fp16 being exact for the grid); elsewhere the integer
     bit-build. Only the taken arm compiles."""
     if _E2M1_HW_CVT:
@@ -288,7 +286,7 @@ def e2m1_cols_to_f32(packed):
 def e2m1_cols_to_e4m3(packed):
     """Column-unpack packed E2M1 (two nibbles per byte along the last dim, low nibble
     first) to E4M3: ``(..., C) uint8 -> (..., 2C)`` — the column-axis counterpart of the
-    row-doubling ``e2m1_to_e4m3``; lossless (every E2M1 value is exact in E4M3). sm_100
     decodes through the CVT unit (``_e2m1_pair_to_e4m3``), elsewhere the integer bit-build;
     either way the bytes are bitcast once. Only the taken arm compiles."""
     if _E2M1_HW_CVT:

 from .bayesian_autotuner import bayesian_autotune
 from .formats import global_scale_stride, is_per_expert_global
+from .compat import FP8_DTYPE, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, compile_time_only_triton_wrap, device_context, is_blackwell, pdl_launch_kwargs
 from .swizzle import swizzle_store_block
 from .scheduling import build_tile_layout, resolve_tile_inline
     return a_fp8, a_scale
+# Hardware FP4 conversions exist on every Blackwell family, independently of TMEM MMA.
+_E2M1_HW_CVT = tl.constexpr(is_blackwell())
 @triton.jit
 def _quant_e2m1_packed(v, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):
     """Pack signed, pre-scaled values on the E2M1 grid to ``(M, K//2)`` uint8 (first value
+    of each pair in the low nibble). On Blackwell the hardware convert
     ``cvt.rn.satfinite.e2m1x2.f32`` (two f32 → one packed byte, first operand → HIGH nibble)
     does it in one instruction; elsewhere a ``>=``-threshold bucketize builds the code. The
     two agree except at exact E2M1 midpoints (0.25, 0.75, …), where the hardware rounds to
 @triton.jit
 def _e2m1_pair_to_f16(packed):
     """One packed E2M1 byte -> the two fp16 values it holds, as a uint32 (low half = low nibble),
+    through the Blackwell CVT unit: ONE instruction per pair against the bit-twiddle form's ~10
     integer ops. Every E2M1 magnitude is exact in fp16, so this is a value-identical decode, and
     the direction of the unpack is irrelevant to the cost — the arithmetic is the whole bill
     (measured on the gpt-oss W4A16 decode weight stream: 39.6us bit-twiddle, either axis, vs
     interleave along K with a SINGLE 3D axis-swap (``join`` -> ``[R, C, 2]`` -> swap last two ->
     reshape) instead of the three 2D transposes ``e2m1_to_e4m3`` needs — the weight-only unpack hot path
     (fp4 weight upcast in-loop, the caller widening to the activation dtype). K order low-nibble-first, matching
+    ``e2m1_to_e4m3``: ``[byte0_lo, byte0_hi, byte1_lo, ...]``. On Blackwell the codes decode through
     the hardware CVT (``_e2m1_pair_to_f16``, fp16 being exact for the grid); elsewhere the integer
     bit-build. Only the taken arm compiles."""
     if _E2M1_HW_CVT:
 def e2m1_cols_to_e4m3(packed):
     """Column-unpack packed E2M1 (two nibbles per byte along the last dim, low nibble
     first) to E4M3: ``(..., C) uint8 -> (..., 2C)`` — the column-axis counterpart of the
+    row-doubling ``e2m1_to_e4m3``; lossless (every E2M1 value is exact in E4M3). Blackwell
     decodes through the CVT unit (``_e2m1_pair_to_e4m3``), elsewhere the integer bit-build;
     either way the bytes are bitcast once. Only the taken arm compiles."""
     if _E2M1_HW_CVT:

build/torch-cuda/symbols.json CHANGED Viewed

@@ -1,11 +1,11 @@
 {
   "schema_version": 1,
-  "module": "_finegrained_kernels_cuda_43f1581",
   "functions": [
     {
       "name": "matmul_2d",
       "qualname": "matmul_2d",
-      "module": "_finegrained_kernels_cuda_43f1581.matmul",
       "kind": "function",
       "doc": "Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of\n``matmul_grouped``/``matmul_batched``, taking the same ``As``/``Bs`` scale spec,\n``activation_format`` / gate|up fusion kwargs, and inferred quant block (no ``block_size`` argument — it\nfalls out of the weight-scale shape, so the data can't disagree with a parameter).\n\n``As`` is the activation block-scale spec, same as the routed ops: ``None`` → the op quantizes raw\n``A`` dynamically; a per-tensor scalar → static (calibrated) FP8 activation scale; block scales →\npre-quantized ``A``. ``Bs`` is the weight block scale. ``a_global_scale``/``b_global_scale`` are the\ntwo-level NVFP4 per-tensor second-level scales (calibrated ``input_scale`` ``g_a`` and\n``weight_scale_2`` ``g_b``) — provided, never computed; the op folds ``g_a·g_b`` onto the\naccumulator (``a_global_scale`` rides raw ``A`` or a pre-quantized ``As`` alike).\n``output_global_scale`` (NVFP4 requant only) is the NEXT proj's provided ``input_scale``:\nthe fused requant normalizes the GLU intermediate by it, returning ``(C, Cs)`` the down-proj\nconsumes as ``As=Cs, a_global_scale=output_global_scale``.\n\n``activation_format`` is the activations' quantization: ``None`` (default) = the weights' own\nformat, ``\"bf16\"`` = weight-only (no activation quant), or an explicit ``\"fp8\"`` / ``\"mxfp8\"`` /\n``\"mxfp4\"`` / ``\"nvfp4\"`` (``\"mxfp8\"`` on MXFP4 weights = the W4A8 chain; ``check_activation_format``\nhas the per-family matrix). ``gate=True`` reads ``B`` as the ``(2N, K)`` gate|up stack and fuses\nthe ``act_fn`` GLU (``swiglu_alpha``/``swiglu_limit`` its clamped variant), returning the\n``[..., N]`` intermediate; ``quantize_output=True`` (MX weights only) stores that intermediate in\n``activation_format`` and returns ``(C, Cs)`` for the next op.\n\nRoutes by weight dtype + inferred block: ``Bs`` None → full-precision; MX/NVFP4 (int8/E4M3 weights\nwith UE8M0 group-32 or E4M3 group-16 ``Bs``) → ``mx_dynamic_matmul``; tensor-wide FP8 (scalar\n``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``; block FP8 → the static variant when ``As`` is a\nper-tensor scalar, else dynamic.",
       "signature": {
@@ -113,7 +113,7 @@
     {
       "name": "matmul_batched",
       "qualname": "matmul_batched",
-      "module": "_finegrained_kernels_cuda_43f1581.batched",
       "kind": "function",
       "doc": "Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one\nprogram per routed row (``expert_ids`` gives its expert).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3 offline for bd/tensor, inline for MX; ``\"bf16\"`` = weight-only;\n``\"mxfp4\"`` = packed E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue. ``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4\nsecond-level scales ride the separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor,\nweights per-expert ``(E,)``; from ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b``\nonto the accumulator. The activation global ``g_a`` is CALIBRATED (the checkpoint's\n``input_scale``): ``a_global_scale=g_a`` with a raw ``A`` has the op quantize ``A / g_a`` per\nblock, and rides a pre-quantized ``As`` the same way. Under an NVFP4 requant the\nfused requant normalizes the GLU intermediate by the PROVIDED ``output_global_scale`` (the next\nproj's calibrated ``input_scale``) before the block quant and returns ``[C, Cs]``; the down\nconsumes it as ``As=Cs, a_global_scale=output_global_scale``. ``gather_idx``/``scatter_idx`` (each None or a ``(S,)`` map)\naddress the source row of ``A`` / destination row of ``C`` per program — None means row\n``s``; the gather lets the gate_up read unexpanded activations (source row\n``s // num_top_k``) with no copy.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_batched`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` → ``mx_dynamic_matmul_batched``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_batched``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``.",
       "signature": {
@@ -239,7 +239,7 @@
     {
       "name": "matmul_grouped",
       "qualname": "matmul_grouped",
-      "module": "_finegrained_kernels_cuda_43f1581.grouped",
       "kind": "function",
       "doc": "Grouped matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision).\n``expert_start`` is the ``(E+1,)`` tiling schedule from one ``compute_grouped_scheduling``\npass, shared by every grouped GEMM of the layer (the expert sort is virtual — nothing is\nphysically permuted).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3, or the fp4 chains; ``\"bf16\"`` = weight-only; ``\"mxfp4\"`` =\npacked E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue.\n``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4 second-level scales ride\nthe separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor, weights per-expert ``(E,)``;\nfrom ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b`` onto the accumulator. The\nactivation global ``g_a`` is CALIBRATED (the checkpoint's ``input_scale``): ``a_global_scale=g_a``\nwith a raw ``A`` has the op quantize ``A / g_a`` per block, and rides a pre-quantized ``As`` the\nsame way. Under an NVFP4 requant the fused requant normalizes the GLU intermediate by the\nPROVIDED ``output_global_scale`` (the next proj's calibrated ``input_scale``) before the block\nquant and returns ``[C, Cs]``; the down consumes it as ``As=Cs, a_global_scale=output_global_scale``.\nRow order is carried by the standalone maps: ``gather_idx`` gathers ``A`` (``None`` ->\nalready expert-ordered), ``scatter_idx`` scatters the output. The fused MoE chain is one\nscheduling pass: gate_up with ``scatter_idx=None`` + ``gate=True`` + ``quantize_output=True``,\nthen down with ``gather_idx=None`` and the\nintermediate's scales as ``As``. EP-sentinel routes sit past ``expert_start[-1]``, which\nno tile reaches.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_grouped`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` → ``mx_dynamic_matmul_grouped``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_grouped``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_grouped``.",
       "signature": {
@@ -365,7 +365,7 @@
     {
       "name": "moe_fused_batched",
       "qualname": "moe_fused_batched",
-      "module": "_finegrained_kernels_cuda_43f1581.moe",
       "kind": "function",
       "doc": "Fused batched MoE (decode): gate_up + SiLU + requant epilogue → per-row quantized\nintermediate → batched down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4 — decode runs the software/swap arms\nbelow the native mxf4nvf4 M=128 staging); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"`` runs the all-fp4 W4A4 chain, ``\"bf16\"`` is weight-only); ``None``\nfollows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each\nstep through the activation dtype so the output matches the unfused reference to\nreduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.",
       "signature": {
@@ -521,7 +521,7 @@
     {
       "name": "moe_fused_grouped",
       "qualname": "moe_fused_grouped",
-      "module": "_finegrained_kernels_cuda_43f1581.moe",
       "kind": "function",
       "doc": "Fused grouped MoE (prefill): gather gate_up + SiLU + requant epilogue → quantized\nexpert-ordered intermediate → grouped down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"``/``\"nvfp4\"`` run all-fp4 W4A4 chains, ``\"bf16\"`` is weight-only);\n``None`` follows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each step through\nthe activation dtype so the output matches the unfused reference to reduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.",
       "signature": {
@@ -674,10 +674,64 @@
         "return_annotation": "torch.Tensor"
       }
     },
     {
       "name": "get_supported_act_fns",
       "qualname": "get_supported_act_fns",
-      "module": "_finegrained_kernels_cuda_43f1581.formats",
       "kind": "function",
       "doc": "The activation names the gate|up epilogue fuses in-kernel; any other activation is applied\non the host by the MoE forwards (``act_fn`` as a callable).",
       "signature": {
@@ -688,7 +742,7 @@
     {
       "name": "get_supported_norms",
       "qualname": "get_supported_norms",
-      "module": "_finegrained_kernels_cuda_43f1581.norm",
       "kind": "function",
       "doc": "The norm names the kernels fuse (``post_expert_norm``), each ``x * factor * rsqrt(mean +\neps)`` over the row: ``\"rms_norm\"`` scales the normalized row by ``weight``,\n``\"centered_rms_norm\"`` by ``1 + weight`` (the zero-centered parameterization), and\n``\"input_scaled_rms_norm\"`` scales by ``1 + weight`` BEFORE normalizing, so the mean square\nis taken on the scaled row. Anything else stays a host callable.",
       "signature": {
@@ -699,7 +753,7 @@
     {
       "name": "rms_norm_rows",
       "qualname": "rms_norm_rows",
-      "module": "_finegrained_kernels_cuda_43f1581.norm",
       "kind": "function",
       "doc": "RMS-normalize ``(S, H)`` rows in one launch, in the ``get_supported_norms()`` form named by\n``norm`` (fp32 accumulate, input dtype out). The MoE chain folds this into\n``weighted_reduce`` instead; this is the standalone form.",
       "signature": {
@@ -735,7 +789,7 @@
     {
       "name": "swizzle_mx_scales",
       "qualname": "swizzle_mx_scales",
-      "module": "_finegrained_kernels_cuda_43f1581.swizzle",
       "kind": "function",
       "doc": "Reorder a block-scale tensor into the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05\nscaled-MMA consumes, one triton launch per matrix (``_swizzle_scales_kernel`` — no torch\npermute/transpose copies). The scale values are unchanged — this only rearranges them from\nplain row-major into the swizzled order the tensor core reads directly (the same layout\ncuBLAS/CUTLASS require for MXFP8/NVFP4 ``scaled_mm``). Plain row-major forces a gather that\ncaps the scaled dot below the fp8/fp4 peak; the swizzle removes it. Apply to weight scales\nonce at model load; the ops accept the result directly.\n\n``scale`` is the per-block scale grid — UE8M0 (``float8_e8m0fnu``, group-32 MXFP8/MXFP4) or\nE4M3 (``float8_e4m3fn``, group-16 NVFP4); the reorder is dtype-agnostic (moves bytes):\n\n- 2D ``(rows, K // group)``: one matrix (a dense weight's scale).\n- 3D ``(E, rows, K // group)``: an expert stack, each expert swizzled independently so its\n  blocks stay ``ceil(rows/128)``-aligned (the reader indexes expert ``e`` at\n  ``e * ceil(rows/128)``; byte-identical to a flat swizzle when ``rows % 128 == 0``).\n\n``gather_idx`` (2D only): a 1-D ``(padded_rows,)`` index mapping each output (sorted) row to\nits source row in ``scale`` (``-1`` = padding → zero row), folded into the kernel's load — the\nrouted/expert-sorted, per-tile-padded layout a grouped GEMM reads affine per BM=128 tile.\n``padded_rows`` must be a multiple of 128.\n\n``rows``/``cols`` are zero-padded to (128, 4) multiples; returns the 5D\n``(G, row_blocks, ceil(cols/4), 2, 256)`` artifact the ops read — ``G`` the expert count (1 for\na matrix), so an expert stack shards/gathers on its leading dim. Bit-identical to CUTLASS's\npacker (verified).",
       "signature": {
@@ -759,7 +813,7 @@
     {
       "name": "unswizzle_mx_scales",
       "qualname": "unswizzle_mx_scales",
-      "module": "_finegrained_kernels_cuda_43f1581.swizzle",
       "kind": "function",
       "doc": "Inverse of ``swizzle_mx_scales``: a SWIZZLE_32_4_4 artifact back to the plain row-major\n``(rows, cols)`` grid — or ``(num_experts, rows, cols)`` for an expert stack — with the\n(128, 4) padding trimmed. Pure reshape/permute on the byte view (no kernel); it runs at\nsave/dequantize time, never on a forward, so the torch-level cost is irrelevant.\n\n``rows``/``cols`` are the LOGICAL dims of ONE matrix (for a gate|up slab, the full ``2N``);\nboth are trimmed back out of the (128, 4) padding. ``num_experts`` selects the stacked\n``(E, rows, cols)`` result and must match the artifact's leading dim.",
       "signature": {
@@ -795,7 +849,7 @@
     {
       "name": "mxfp8_act_quant",
       "qualname": "mxfp8_act_quant",
-      "module": "_finegrained_kernels_cuda_43f1581.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MX once (E4M3 values + UE8M0 group-32 uint8 scales)\ninstead of inline per weight-tile — the fused gate_up re-ran the inline quant per N-tile\n(16x redundant amax/convert ALU + 2x act bytes), which held it at ~380 TFLOPS while the\npre-quantized down kernel ran at ~1080. One pass costs ~50µs at 8k tokens. Bit-exact with\nthe inline form (same group boundaries). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path (same dense grid, per-element store).",
       "signature": {
@@ -819,7 +873,7 @@
     {
       "name": "mxfp4_act_quant",
       "qualname": "mxfp4_act_quant",
-      "module": "_finegrained_kernels_cuda_43f1581.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MXFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8, two codes per byte, first value in the low nibble) + UE8M0\ngroup-32 uint8 scales (``(T, K//32)``, amax/6 ceil'd to a power of two). Bit-identical\nto the fused epilogues' inline form (shared ``mx_act_quant_inline`` arm). Feeds the\nW4A4 arm of the MX matmul ops; quantizing activations to fp4 at runtime is an\naccuracy call the caller owns — the ops never do it implicitly. ``swizzled=True`` emits the\nscale directly in SWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {
@@ -843,7 +897,7 @@
     {
       "name": "nvfp4_act_quant",
       "qualname": "nvfp4_act_quant",
-      "module": "_finegrained_kernels_cuda_43f1581.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to NVFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8) + E4M3 group-16 block scales (``(T, K//16)`` — ``amax/6`` rounded to\nE4M3, NOT a power of two; values divide by the DECODED scale before the E2M1 grid,\nthe standard two-step). ``global_scale`` is the CALIBRATED per-tensor second level\n(fp32, the checkpoint's ``input_scale``): values are normalized by it before\nthe block quant, so block scales stay in e4m3 range regardless of the activation's\ndynamic range — the canonical two-level format. The GEMM folds ``g_a·g_b`` back onto\nthe accumulator (pass ``As=scales`` with ``a_global_scale=global_scale``). Per-expert globals\nneed each row's expert, which ``expert_index`` gives where the caller already holds that map\n(the batched op's ``expert_ids``, sentinels and all) — it is read only when the globals are\nper expert, and ``mx_act_quant_grouped`` needs none at all, its tiles carrying an expert.\n``None`` = single-level (``g_a = 1``). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {

 {
   "schema_version": 1,
+  "module": "_finegrained_kernels_cuda_ec5cc90",
   "functions": [
     {
       "name": "matmul_2d",
       "qualname": "matmul_2d",
+      "module": "_finegrained_kernels_cuda_ec5cc90.matmul",
       "kind": "function",
       "doc": "Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of\n``matmul_grouped``/``matmul_batched``, taking the same ``As``/``Bs`` scale spec,\n``activation_format`` / gate|up fusion kwargs, and inferred quant block (no ``block_size`` argument — it\nfalls out of the weight-scale shape, so the data can't disagree with a parameter).\n\n``As`` is the activation block-scale spec, same as the routed ops: ``None`` → the op quantizes raw\n``A`` dynamically; a per-tensor scalar → static (calibrated) FP8 activation scale; block scales →\npre-quantized ``A``. ``Bs`` is the weight block scale. ``a_global_scale``/``b_global_scale`` are the\ntwo-level NVFP4 per-tensor second-level scales (calibrated ``input_scale`` ``g_a`` and\n``weight_scale_2`` ``g_b``) — provided, never computed; the op folds ``g_a·g_b`` onto the\naccumulator (``a_global_scale`` rides raw ``A`` or a pre-quantized ``As`` alike).\n``output_global_scale`` (NVFP4 requant only) is the NEXT proj's provided ``input_scale``:\nthe fused requant normalizes the GLU intermediate by it, returning ``(C, Cs)`` the down-proj\nconsumes as ``As=Cs, a_global_scale=output_global_scale``.\n\n``activation_format`` is the activations' quantization: ``None`` (default) = the weights' own\nformat, ``\"bf16\"`` = weight-only (no activation quant), or an explicit ``\"fp8\"`` / ``\"mxfp8\"`` /\n``\"mxfp4\"`` / ``\"nvfp4\"`` (``\"mxfp8\"`` on MXFP4 weights = the W4A8 chain; ``check_activation_format``\nhas the per-family matrix). ``gate=True`` reads ``B`` as the ``(2N, K)`` gate|up stack and fuses\nthe ``act_fn`` GLU (``swiglu_alpha``/``swiglu_limit`` its clamped variant), returning the\n``[..., N]`` intermediate; ``quantize_output=True`` (MX weights only) stores that intermediate in\n``activation_format`` and returns ``(C, Cs)`` for the next op.\n\nRoutes by weight dtype + inferred block: ``Bs`` None → full-precision; MX/NVFP4 (int8/E4M3 weights\nwith UE8M0 group-32 or E4M3 group-16 ``Bs``) → ``mx_dynamic_matmul``; tensor-wide FP8 (scalar\n``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``; block FP8 → the static variant when ``As`` is a\nper-tensor scalar, else dynamic.",
       "signature": {
     {
       "name": "matmul_batched",
       "qualname": "matmul_batched",
+      "module": "_finegrained_kernels_cuda_ec5cc90.batched",
       "kind": "function",
       "doc": "Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one\nprogram per routed row (``expert_ids`` gives its expert).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3 offline for bd/tensor, inline for MX; ``\"bf16\"`` = weight-only;\n``\"mxfp4\"`` = packed E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue. ``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4\nsecond-level scales ride the separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor,\nweights per-expert ``(E,)``; from ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b``\nonto the accumulator. The activation global ``g_a`` is CALIBRATED (the checkpoint's\n``input_scale``): ``a_global_scale=g_a`` with a raw ``A`` has the op quantize ``A / g_a`` per\nblock, and rides a pre-quantized ``As`` the same way. Under an NVFP4 requant the\nfused requant normalizes the GLU intermediate by the PROVIDED ``output_global_scale`` (the next\nproj's calibrated ``input_scale``) before the block quant and returns ``[C, Cs]``; the down\nconsumes it as ``As=Cs, a_global_scale=output_global_scale``. ``gather_idx``/``scatter_idx`` (each None or a ``(S,)`` map)\naddress the source row of ``A`` / destination row of ``C`` per program — None means row\n``s``; the gather lets the gate_up read unexpanded activations (source row\n``s // num_top_k``) with no copy.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_batched`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` → ``mx_dynamic_matmul_batched``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_batched``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``.",
       "signature": {
     {
       "name": "matmul_grouped",
       "qualname": "matmul_grouped",
+      "module": "_finegrained_kernels_cuda_ec5cc90.grouped",
       "kind": "function",
       "doc": "Grouped matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision).\n``expert_start`` is the ``(E+1,)`` tiling schedule from one ``compute_grouped_scheduling``\npass, shared by every grouped GEMM of the layer (the expert sort is virtual — nothing is\nphysically permuted).\n\n``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized\nintermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static\n(calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against\nit; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =\nthe weights' own format — fp8/E4M3, or the fp4 chains; ``\"bf16\"`` = weight-only; ``\"mxfp4\"`` =\npacked E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.\n``quantize_output`` requantizes the output into ``activation_format`` — the return is then\n``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the\nepilogue.\n``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4 second-level scales ride\nthe separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor, weights per-expert ``(E,)``;\nfrom ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b`` onto the accumulator. The\nactivation global ``g_a`` is CALIBRATED (the checkpoint's ``input_scale``): ``a_global_scale=g_a``\nwith a raw ``A`` has the op quantize ``A / g_a`` per block, and rides a pre-quantized ``As`` the\nsame way. Under an NVFP4 requant the fused requant normalizes the GLU intermediate by the\nPROVIDED ``output_global_scale`` (the next proj's calibrated ``input_scale``) before the block\nquant and returns ``[C, Cs]``; the down consumes it as ``As=Cs, a_global_scale=output_global_scale``.\nRow order is carried by the standalone maps: ``gather_idx`` gathers ``A`` (``None`` ->\nalready expert-ordered), ``scatter_idx`` scatters the output. The fused MoE chain is one\nscheduling pass: gate_up with ``scatter_idx=None`` + ``gate=True`` + ``quantize_output=True``,\nthen down with ``gather_idx=None`` and the\nintermediate's scales as ``As``. EP-sentinel routes sit past ``expert_start[-1]``, which\nno tile reaches.\n\nRoutes by what the weight tensors themselves say (there is no ``block_size``\nparameter — the quantization block is derived from the scale shape,\n``weight_block_size``):\n- ``Bs`` None → ``full_precision_matmul_grouped`` (plain dot, no scales anywhere).\n- MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0\n  group-32 ``Bs`` → ``mx_dynamic_matmul_grouped``.\n- one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →\n  ``w8a8_tensor_dynamic_fp8_matmul_grouped``.\n- block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_grouped``.",
       "signature": {
     {
       "name": "moe_fused_batched",
       "qualname": "moe_fused_batched",
+      "module": "_finegrained_kernels_cuda_ec5cc90.moe",
       "kind": "function",
       "doc": "Fused batched MoE (decode): gate_up + SiLU + requant epilogue → per-row quantized\nintermediate → batched down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4 — decode runs the software/swap arms\nbelow the native mxf4nvf4 M=128 staging); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"`` runs the all-fp4 W4A4 chain, ``\"bf16\"`` is weight-only); ``None``\nfollows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each\nstep through the activation dtype so the output matches the unfused reference to\nreduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.",
       "signature": {
     {
       "name": "moe_fused_grouped",
       "qualname": "moe_fused_grouped",
+      "module": "_finegrained_kernels_cuda_ec5cc90.moe",
       "kind": "function",
       "doc": "Fused grouped MoE (prefill): gather gate_up + SiLU + requant epilogue → quantized\nexpert-ordered intermediate → grouped down → routing-weighted top-k reduce. Returns\n``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale\nlayout (block-dynamic FP8, MXFP8/MXFP4, NVFP4); ``activation_format`` names the activation\nquantization for the whole block — activations and the fused intermediate requant\ncarry it (``\"mxfp4\"``/``\"nvfp4\"`` run all-fp4 W4A4 chains, ``\"bf16\"`` is weight-only);\n``None`` follows the weight format, and the ops validate the pairing. The\n``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's\ncalibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,\nrequants the intermediate against the down's, and the down consumes that as its activation\nglobal. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each step through\nthe activation dtype so the output matches the unfused reference to reduce order. ``act_fn`` is a\n``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any\ncallable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated\nprojection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)\n— the swizzled layout is for the dot_scaled arm; weight-only chains take affine.",
       "signature": {
         "return_annotation": "torch.Tensor"
       }
     },
+    {
+      "name": "embedding_lookup",
+      "qualname": "embedding_lookup",
+      "module": "_finegrained_kernels_cuda_ec5cc90.embedding",
+      "kind": "function",
+      "doc": "Gather and dequantize selected rows without expanding the embedding table.\n\n``weight_format`` selects ``fp8``, ``nvfp4``, ``mxfp8``, ``mxfp4``, ``int8``, or ``int4``.\nBy default, a global scale selects NVFP4, signed int8 weights select INT8, and otherwise\nFP8 is used. FP8 accepts scalar, per-row, or square block scales (``block_size``).\nNVFP4 uses E4M3 group-16 scales and one FP32 multiplier; MX formats use UE8M0 group-32\nscales. FP4 packs adjacent columns, low nibble first. INT8 uses per-row scales; INT4\npacks adjacent vocabulary rows and uses per-row column groups inferred from its scale\ngrid. Integer values are signed, symmetric, and have no zero point.\n\nOutputs default to the floating scale dtype, or bf16 for NVFP4/MX. Table/grid strides\nmay describe TP slices. Indices follow ``torch.nn.functional.embedding``: negative and\nout-of-range values fail; padding indices look up their stored row. Empty, repeated,\nand noncontiguous indices work.",
+      "signature": {
+        "parameters": [
+          {
+            "name": "weight",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "indices",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": null,
+            "annotation": "torch.Tensor"
+          },
+          {
+            "name": "global_scale",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.Tensor | None"
+          },
+          {
+            "name": "output_dtype",
+            "kind": "POSITIONAL_OR_KEYWORD",
+            "default": "None",
+            "annotation": "torch.dtype | None"
+          },
+          {
+            "name": "weight_format",
+            "kind": "KEYWORD_ONLY",
+            "default": "None",
+            "annotation": "str | None"
+          },
+          {
+            "name": "block_size",
+            "kind": "KEYWORD_ONLY",
+            "default": "128",
+            "annotation": "int"
+          }
+        ],
+        "return_annotation": "torch.Tensor"
+      }
+    },
     {
       "name": "get_supported_act_fns",
       "qualname": "get_supported_act_fns",
+      "module": "_finegrained_kernels_cuda_ec5cc90.formats",
       "kind": "function",
       "doc": "The activation names the gate|up epilogue fuses in-kernel; any other activation is applied\non the host by the MoE forwards (``act_fn`` as a callable).",
       "signature": {
     {
       "name": "get_supported_norms",
       "qualname": "get_supported_norms",
+      "module": "_finegrained_kernels_cuda_ec5cc90.norm",
       "kind": "function",
       "doc": "The norm names the kernels fuse (``post_expert_norm``), each ``x * factor * rsqrt(mean +\neps)`` over the row: ``\"rms_norm\"`` scales the normalized row by ``weight``,\n``\"centered_rms_norm\"`` by ``1 + weight`` (the zero-centered parameterization), and\n``\"input_scaled_rms_norm\"`` scales by ``1 + weight`` BEFORE normalizing, so the mean square\nis taken on the scaled row. Anything else stays a host callable.",
       "signature": {
     {
       "name": "rms_norm_rows",
       "qualname": "rms_norm_rows",
+      "module": "_finegrained_kernels_cuda_ec5cc90.norm",
       "kind": "function",
       "doc": "RMS-normalize ``(S, H)`` rows in one launch, in the ``get_supported_norms()`` form named by\n``norm`` (fp32 accumulate, input dtype out). The MoE chain folds this into\n``weighted_reduce`` instead; this is the standalone form.",
       "signature": {
     {
       "name": "swizzle_mx_scales",
       "qualname": "swizzle_mx_scales",
+      "module": "_finegrained_kernels_cuda_ec5cc90.swizzle",
       "kind": "function",
       "doc": "Reorder a block-scale tensor into the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05\nscaled-MMA consumes, one triton launch per matrix (``_swizzle_scales_kernel`` — no torch\npermute/transpose copies). The scale values are unchanged — this only rearranges them from\nplain row-major into the swizzled order the tensor core reads directly (the same layout\ncuBLAS/CUTLASS require for MXFP8/NVFP4 ``scaled_mm``). Plain row-major forces a gather that\ncaps the scaled dot below the fp8/fp4 peak; the swizzle removes it. Apply to weight scales\nonce at model load; the ops accept the result directly.\n\n``scale`` is the per-block scale grid — UE8M0 (``float8_e8m0fnu``, group-32 MXFP8/MXFP4) or\nE4M3 (``float8_e4m3fn``, group-16 NVFP4); the reorder is dtype-agnostic (moves bytes):\n\n- 2D ``(rows, K // group)``: one matrix (a dense weight's scale).\n- 3D ``(E, rows, K // group)``: an expert stack, each expert swizzled independently so its\n  blocks stay ``ceil(rows/128)``-aligned (the reader indexes expert ``e`` at\n  ``e * ceil(rows/128)``; byte-identical to a flat swizzle when ``rows % 128 == 0``).\n\n``gather_idx`` (2D only): a 1-D ``(padded_rows,)`` index mapping each output (sorted) row to\nits source row in ``scale`` (``-1`` = padding → zero row), folded into the kernel's load — the\nrouted/expert-sorted, per-tile-padded layout a grouped GEMM reads affine per BM=128 tile.\n``padded_rows`` must be a multiple of 128.\n\n``rows``/``cols`` are zero-padded to (128, 4) multiples; returns the 5D\n``(G, row_blocks, ceil(cols/4), 2, 256)`` artifact the ops read — ``G`` the expert count (1 for\na matrix), so an expert stack shards/gathers on its leading dim. Bit-identical to CUTLASS's\npacker (verified).",
       "signature": {
     {
       "name": "unswizzle_mx_scales",
       "qualname": "unswizzle_mx_scales",
+      "module": "_finegrained_kernels_cuda_ec5cc90.swizzle",
       "kind": "function",
       "doc": "Inverse of ``swizzle_mx_scales``: a SWIZZLE_32_4_4 artifact back to the plain row-major\n``(rows, cols)`` grid — or ``(num_experts, rows, cols)`` for an expert stack — with the\n(128, 4) padding trimmed. Pure reshape/permute on the byte view (no kernel); it runs at\nsave/dequantize time, never on a forward, so the torch-level cost is irrelevant.\n\n``rows``/``cols`` are the LOGICAL dims of ONE matrix (for a gate|up slab, the full ``2N``);\nboth are trimmed back out of the (128, 4) padding. ``num_experts`` selects the stacked\n``(E, rows, cols)`` result and must match the artifact's leading dim.",
       "signature": {
     {
       "name": "mxfp8_act_quant",
       "qualname": "mxfp8_act_quant",
+      "module": "_finegrained_kernels_cuda_ec5cc90.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MX once (E4M3 values + UE8M0 group-32 uint8 scales)\ninstead of inline per weight-tile — the fused gate_up re-ran the inline quant per N-tile\n(16x redundant amax/convert ALU + 2x act bytes), which held it at ~380 TFLOPS while the\npre-quantized down kernel ran at ~1080. One pass costs ~50µs at 8k tokens. Bit-exact with\nthe inline form (same group boundaries). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path (same dense grid, per-element store).",
       "signature": {
     {
       "name": "mxfp4_act_quant",
       "qualname": "mxfp4_act_quant",
+      "module": "_finegrained_kernels_cuda_ec5cc90.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to MXFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8, two codes per byte, first value in the low nibble) + UE8M0\ngroup-32 uint8 scales (``(T, K//32)``, amax/6 ceil'd to a power of two). Bit-identical\nto the fused epilogues' inline form (shared ``mx_act_quant_inline`` arm). Feeds the\nW4A4 arm of the MX matmul ops; quantizing activations to fp4 at runtime is an\naccuracy call the caller owns — the ops never do it implicitly. ``swizzled=True`` emits the\nscale directly in SWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {
     {
       "name": "nvfp4_act_quant",
       "qualname": "nvfp4_act_quant",
+      "module": "_finegrained_kernels_cuda_ec5cc90.quant",
       "kind": "function",
       "doc": "Quantize ``(T, K)`` activations to NVFP4 in one kernel pass: packed-E2M1 values\n(``(T, K//2)`` int8) + E4M3 group-16 block scales (``(T, K//16)`` — ``amax/6`` rounded to\nE4M3, NOT a power of two; values divide by the DECODED scale before the E2M1 grid,\nthe standard two-step). ``global_scale`` is the CALIBRATED per-tensor second level\n(fp32, the checkpoint's ``input_scale``): values are normalized by it before\nthe block quant, so block scales stay in e4m3 range regardless of the activation's\ndynamic range — the canonical two-level format. The GEMM folds ``g_a·g_b`` back onto\nthe accumulator (pass ``As=scales`` with ``a_global_scale=global_scale``). Per-expert globals\nneed each row's expert, which ``expert_index`` gives where the caller already holds that map\n(the batched op's ``expert_ids``, sentinels and all) — it is read only when the globals are\nper expert, and ``mx_act_quant_grouped`` needs none at all, its tiles carrying an expert.\n``None`` = single-level (``g_a = 1``). ``swizzled=True`` emits the scale directly in\nSWIZZLE_32_4_4 for the tcgen05 fast path.",
       "signature": {