nbxtensor 0.0.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (291) hide show
  1. nbxtensor/THIRD_PARTY_NOTICES.md +287 -0
  2. nbxtensor/__init__.py +17 -0
  3. nbxtensor/backend/__init__.py +1 -0
  4. nbxtensor/backend/capabilities.py +207 -0
  5. nbxtensor/backend/cpu/__init__.py +1 -0
  6. nbxtensor/backend/cpu/blas.py +102 -0
  7. nbxtensor/backend/cpu/pool.py +204 -0
  8. nbxtensor/backend/cpu_backend.py +162 -0
  9. nbxtensor/backend/device_facts.py +105 -0
  10. nbxtensor/backend/device_files.py +129 -0
  11. nbxtensor/backend/device_names.py +18 -0
  12. nbxtensor/backend/devices/amd/cdna.yml +228 -0
  13. nbxtensor/backend/devices/amd/cdna2.yml +228 -0
  14. nbxtensor/backend/devices/amd/cdna3.yml +227 -0
  15. nbxtensor/backend/devices/apple/apple_m1.yml +266 -0
  16. nbxtensor/backend/devices/apple/apple_m1_max.yml +266 -0
  17. nbxtensor/backend/devices/apple/apple_m1_pro.yml +266 -0
  18. nbxtensor/backend/devices/apple/apple_m1_ultra.yml +266 -0
  19. nbxtensor/backend/devices/apple/apple_m2.yml +266 -0
  20. nbxtensor/backend/devices/apple/apple_m2_max.yml +266 -0
  21. nbxtensor/backend/devices/apple/apple_m2_pro.yml +266 -0
  22. nbxtensor/backend/devices/apple/apple_m2_ultra.yml +266 -0
  23. nbxtensor/backend/devices/apple/apple_m3.yml +266 -0
  24. nbxtensor/backend/devices/apple/apple_m3_max.yml +266 -0
  25. nbxtensor/backend/devices/apple/apple_m3_pro.yml +266 -0
  26. nbxtensor/backend/devices/apple/apple_m3_ultra.yml +266 -0
  27. nbxtensor/backend/devices/apple/apple_m4.yml +266 -0
  28. nbxtensor/backend/devices/apple/apple_m4_max.yml +266 -0
  29. nbxtensor/backend/devices/apple/apple_m4_pro.yml +308 -0
  30. nbxtensor/backend/devices/apple/apple_m5.yml +269 -0
  31. nbxtensor/backend/devices/apple/apple_m5_max.yml +269 -0
  32. nbxtensor/backend/devices/apple/apple_m5_pro.yml +269 -0
  33. nbxtensor/backend/devices/apple/apple_m5_ultra.yml +269 -0
  34. nbxtensor/backend/devices/apple/apple_m6.yml +269 -0
  35. nbxtensor/backend/devices/apple/apple_silicon.yml +263 -0
  36. nbxtensor/backend/devices/nvidia/ampere.yml +214 -0
  37. nbxtensor/backend/devices/nvidia/hopper.yml +213 -0
  38. nbxtensor/backend/devices/nvidia/volta.yml +411 -0
  39. nbxtensor/backend/faults.py +151 -0
  40. nbxtensor/backend/host_memory.py +185 -0
  41. nbxtensor/backend/launcher.py +2350 -0
  42. nbxtensor/backend/launcher_contract.py +429 -0
  43. nbxtensor/backend/metal_backend.py +661 -0
  44. nbxtensor/backend/metal_device.py +1274 -0
  45. nbxtensor/backend/process_memory.py +73 -0
  46. nbxtensor/backend/replay_graph.py +322 -0
  47. nbxtensor/backend/runtime.py +409 -0
  48. nbxtensor/backend/triton_ext_driver.py +735 -0
  49. nbxtensor/dtype/__init__.py +1 -0
  50. nbxtensor/dtype/convert.py +107 -0
  51. nbxtensor/dtype/dtypes.yml +70 -0
  52. nbxtensor/dtype/itemsize.py +166 -0
  53. nbxtensor/dtype/nbx_dtype.py +72 -0
  54. nbxtensor/dtype/operands.py +158 -0
  55. nbxtensor/dtype/outputs.py +64 -0
  56. nbxtensor/dtype/policy.py +155 -0
  57. nbxtensor/dtype/policy.yml +51 -0
  58. nbxtensor/dtype/promotion.py +112 -0
  59. nbxtensor/host.py +55 -0
  60. nbxtensor/io/__init__.py +1 -0
  61. nbxtensor/io/read_rate.py +197 -0
  62. nbxtensor/io/safetensors.py +314 -0
  63. nbxtensor/memory/__init__.py +9 -0
  64. nbxtensor/memory/allocator.py +1988 -0
  65. nbxtensor/memory/arena.py +75 -0
  66. nbxtensor/memory/device_memory.py +150 -0
  67. nbxtensor/memory/device_transfer.py +154 -0
  68. nbxtensor/memory/layout.py +43 -0
  69. nbxtensor/memory/memory_pool.py +191 -0
  70. nbxtensor/memory/metal.py +96 -0
  71. nbxtensor/memory/oom_advice.py +91 -0
  72. nbxtensor/memory/pointer_tables.py +323 -0
  73. nbxtensor/memory/pool.py +507 -0
  74. nbxtensor/memory/shadow.py +43 -0
  75. nbxtensor/memory/slab.py +273 -0
  76. nbxtensor/memory/trace_replay.py +122 -0
  77. nbxtensor/memory/traces/sana-video-2b-720p.16g.2026-11-02.tsv +13610 -0
  78. nbxtensor/memory/vmm.py +118 -0
  79. nbxtensor/ops/__init__.py +1 -0
  80. nbxtensor/ops/_autotune_policy.py +90 -0
  81. nbxtensor/ops/_common.py +32 -0
  82. nbxtensor/ops/_configs.py +1041 -0
  83. nbxtensor/ops/_cost.py +231 -0
  84. nbxtensor/ops/abs.py +17 -0
  85. nbxtensor/ops/adaptive_avg_pool2d.py +161 -0
  86. nbxtensor/ops/add.py +96 -0
  87. nbxtensor/ops/addcdiv.py +29 -0
  88. nbxtensor/ops/addcmul.py +29 -0
  89. nbxtensor/ops/addmv_op.py +64 -0
  90. nbxtensor/ops/addr.py +68 -0
  91. nbxtensor/ops/all_reduce.py +83 -0
  92. nbxtensor/ops/amax.py +40 -0
  93. nbxtensor/ops/any_reduce.py +83 -0
  94. nbxtensor/ops/arange_op.py +27 -0
  95. nbxtensor/ops/argmax.py +107 -0
  96. nbxtensor/ops/argmin.py +93 -0
  97. nbxtensor/ops/atan2.py +71 -0
  98. nbxtensor/ops/attention_wrappers.py +1648 -0
  99. nbxtensor/ops/autotune_bucket.py +137 -0
  100. nbxtensor/ops/autotune_cache.py +430 -0
  101. nbxtensor/ops/autotune_certified.py +1262 -0
  102. nbxtensor/ops/autotune_refusals.py +380 -0
  103. nbxtensor/ops/autotune_settings.py +58 -0
  104. nbxtensor/ops/avg_pool2d.py +73 -0
  105. nbxtensor/ops/baddbmm_op.py +132 -0
  106. nbxtensor/ops/batch_norm.py +163 -0
  107. nbxtensor/ops/binary_strided.py +196 -0
  108. nbxtensor/ops/bitwise_and.py +26 -0
  109. nbxtensor/ops/bitwise_left_shift.py +47 -0
  110. nbxtensor/ops/bitwise_not.py +24 -0
  111. nbxtensor/ops/bitwise_or.py +26 -0
  112. nbxtensor/ops/bitwise_right_shift.py +47 -0
  113. nbxtensor/ops/bucketize.py +37 -0
  114. nbxtensor/ops/cat_op.py +87 -0
  115. nbxtensor/ops/celu.py +20 -0
  116. nbxtensor/ops/clamp.py +24 -0
  117. nbxtensor/ops/clamp_min.py +22 -0
  118. nbxtensor/ops/conv2d.py +265 -0
  119. nbxtensor/ops/conv2d_m8n8k4.py +165 -0
  120. nbxtensor/ops/conv_depthwise2d.py +154 -0
  121. nbxtensor/ops/conv_transpose2d.py +148 -0
  122. nbxtensor/ops/conv_wrappers.py +1482 -0
  123. nbxtensor/ops/copy.py +18 -0
  124. nbxtensor/ops/copy_op.py +39 -0
  125. nbxtensor/ops/cos.py +17 -0
  126. nbxtensor/ops/count_nonzero.py +64 -0
  127. nbxtensor/ops/cross_entropy.py +67 -0
  128. nbxtensor/ops/cummax.py +109 -0
  129. nbxtensor/ops/cummin.py +105 -0
  130. nbxtensor/ops/cumsum.py +244 -0
  131. nbxtensor/ops/decode_attn_vec.py +162 -0
  132. nbxtensor/ops/depthwise_conv2d.py +132 -0
  133. nbxtensor/ops/dequant_gemv.py +190 -0
  134. nbxtensor/ops/dft.py +67 -0
  135. nbxtensor/ops/dispatch.py +1228 -0
  136. nbxtensor/ops/div.py +75 -0
  137. nbxtensor/ops/dot_op.py +67 -0
  138. nbxtensor/ops/dropout.py +23 -0
  139. nbxtensor/ops/dtype_convert.py +42 -0
  140. nbxtensor/ops/elementwise_launch.py +81 -0
  141. nbxtensor/ops/elementwise_wrappers.py +2315 -0
  142. nbxtensor/ops/elu.py +20 -0
  143. nbxtensor/ops/embedding.py +69 -0
  144. nbxtensor/ops/eq.py +32 -0
  145. nbxtensor/ops/erf.py +18 -0
  146. nbxtensor/ops/exp.py +17 -0
  147. nbxtensor/ops/exp2.py +18 -0
  148. nbxtensor/ops/fft_op.py +125 -0
  149. nbxtensor/ops/fill_op.py +20 -0
  150. nbxtensor/ops/flash_attention.py +307 -0
  151. nbxtensor/ops/flash_attention_m8n8k4.py +384 -0
  152. nbxtensor/ops/flash_decode.py +78 -0
  153. nbxtensor/ops/flip_op.py +60 -0
  154. nbxtensor/ops/floor.py +90 -0
  155. nbxtensor/ops/full_op.py +24 -0
  156. nbxtensor/ops/fused_moe.py +417 -0
  157. nbxtensor/ops/fused_moe_m8n8k4.py +99 -0
  158. nbxtensor/ops/fused_upsample_conv.py +413 -0
  159. nbxtensor/ops/gather_op.py +52 -0
  160. nbxtensor/ops/ge.py +32 -0
  161. nbxtensor/ops/gelu.py +31 -0
  162. nbxtensor/ops/gemv_vec.py +80 -0
  163. nbxtensor/ops/glu.py +43 -0
  164. nbxtensor/ops/grid_sampler.py +145 -0
  165. nbxtensor/ops/group_norm_tiled.py +67 -0
  166. nbxtensor/ops/groupnorm.py +123 -0
  167. nbxtensor/ops/gt.py +32 -0
  168. nbxtensor/ops/hardsigmoid.py +18 -0
  169. nbxtensor/ops/hardswish.py +19 -0
  170. nbxtensor/ops/im2col.py +57 -0
  171. nbxtensor/ops/index_add.py +95 -0
  172. nbxtensor/ops/index_put_op.py +92 -0
  173. nbxtensor/ops/index_select.py +121 -0
  174. nbxtensor/ops/index_wrappers.py +770 -0
  175. nbxtensor/ops/isfinite.py +20 -0
  176. nbxtensor/ops/isinf.py +20 -0
  177. nbxtensor/ops/isnan.py +18 -0
  178. nbxtensor/ops/launch_keys.py +1054 -0
  179. nbxtensor/ops/launch_observer.py +54 -0
  180. nbxtensor/ops/layernorm.py +71 -0
  181. nbxtensor/ops/le.py +35 -0
  182. nbxtensor/ops/leaky_relu.py +19 -0
  183. nbxtensor/ops/lerp.py +76 -0
  184. nbxtensor/ops/linalg_wrappers.py +1370 -0
  185. nbxtensor/ops/linspace_op.py +47 -0
  186. nbxtensor/ops/log.py +17 -0
  187. nbxtensor/ops/log_sigmoid.py +23 -0
  188. nbxtensor/ops/logical_and.py +26 -0
  189. nbxtensor/ops/logical_not.py +24 -0
  190. nbxtensor/ops/logical_or.py +26 -0
  191. nbxtensor/ops/logical_xor.py +37 -0
  192. nbxtensor/ops/loss_wrappers.py +282 -0
  193. nbxtensor/ops/lt.py +32 -0
  194. nbxtensor/ops/masked_fill.py +21 -0
  195. nbxtensor/ops/matmul.py +377 -0
  196. nbxtensor/ops/matmul_m8n8k4.py +235 -0
  197. nbxtensor/ops/max_pool2d.py +74 -0
  198. nbxtensor/ops/max_pool3d.py +93 -0
  199. nbxtensor/ops/max_reduce.py +91 -0
  200. nbxtensor/ops/maximum.py +18 -0
  201. nbxtensor/ops/mean.py +40 -0
  202. nbxtensor/ops/min_reduce.py +91 -0
  203. nbxtensor/ops/minimum.py +18 -0
  204. nbxtensor/ops/mish.py +22 -0
  205. nbxtensor/ops/mma_m8n8k4.py +152 -0
  206. nbxtensor/ops/moe_align.py +195 -0
  207. nbxtensor/ops/moe_decode_vec.py +374 -0
  208. nbxtensor/ops/moe_device_residency.py +141 -0
  209. nbxtensor/ops/moe_wrappers.py +345 -0
  210. nbxtensor/ops/mse_loss.py +79 -0
  211. nbxtensor/ops/mul.py +51 -0
  212. nbxtensor/ops/nan_to_num.py +23 -0
  213. nbxtensor/ops/ne.py +32 -0
  214. nbxtensor/ops/neg.py +17 -0
  215. nbxtensor/ops/nllloss.py +124 -0
  216. nbxtensor/ops/nonzero.py +33 -0
  217. nbxtensor/ops/nonzero_op.py +62 -0
  218. nbxtensor/ops/norm_wrappers.py +579 -0
  219. nbxtensor/ops/ones_op.py +20 -0
  220. nbxtensor/ops/oracles/__init__.py +6 -0
  221. nbxtensor/ops/oracles/conv2d_fp64.py +450 -0
  222. nbxtensor/ops/pad_op.py +159 -0
  223. nbxtensor/ops/paged_decode_attn.py +215 -0
  224. nbxtensor/ops/pixel_shuffle.py +144 -0
  225. nbxtensor/ops/pixel_unshuffle.py +72 -0
  226. nbxtensor/ops/pow.py +137 -0
  227. nbxtensor/ops/prod.py +81 -0
  228. nbxtensor/ops/rand_op.py +52 -0
  229. nbxtensor/ops/random_wrappers.py +366 -0
  230. nbxtensor/ops/reciprocal.py +17 -0
  231. nbxtensor/ops/reduction_wrappers.py +845 -0
  232. nbxtensor/ops/relu.py +18 -0
  233. nbxtensor/ops/remainder.py +46 -0
  234. nbxtensor/ops/repeat_interleave.py +48 -0
  235. nbxtensor/ops/repeat_op.py +132 -0
  236. nbxtensor/ops/residual_chain.py +374 -0
  237. nbxtensor/ops/rmsnorm.py +68 -0
  238. nbxtensor/ops/rope.py +142 -0
  239. nbxtensor/ops/rsqrt.py +17 -0
  240. nbxtensor/ops/scaled_softmax.py +96 -0
  241. nbxtensor/ops/scatter_op.py +139 -0
  242. nbxtensor/ops/screen_oracle.py +479 -0
  243. nbxtensor/ops/selu.py +22 -0
  244. nbxtensor/ops/shape_wrappers.py +1032 -0
  245. nbxtensor/ops/sigmoid.py +19 -0
  246. nbxtensor/ops/silu.py +27 -0
  247. nbxtensor/ops/sin.py +17 -0
  248. nbxtensor/ops/softmax.py +169 -0
  249. nbxtensor/ops/softplus.py +27 -0
  250. nbxtensor/ops/sort_op.py +304 -0
  251. nbxtensor/ops/spectral_wrappers.py +481 -0
  252. nbxtensor/ops/sqrt.py +17 -0
  253. nbxtensor/ops/stack_op.py +84 -0
  254. nbxtensor/ops/std.py +113 -0
  255. nbxtensor/ops/strided_copy.py +200 -0
  256. nbxtensor/ops/sub.py +35 -0
  257. nbxtensor/ops/sum.py +40 -0
  258. nbxtensor/ops/tan.py +19 -0
  259. nbxtensor/ops/tanh.py +21 -0
  260. nbxtensor/ops/threshold.py +19 -0
  261. nbxtensor/ops/tiling.yml +91 -0
  262. nbxtensor/ops/tiling_policy.py +43 -0
  263. nbxtensor/ops/tiling_sizes.py +370 -0
  264. nbxtensor/ops/topk.py +267 -0
  265. nbxtensor/ops/trace_op.py +64 -0
  266. nbxtensor/ops/tril.py +64 -0
  267. nbxtensor/ops/triu.py +62 -0
  268. nbxtensor/ops/unfold_op.py +37 -0
  269. nbxtensor/ops/unique_op.py +97 -0
  270. nbxtensor/ops/upsample_bilinear2d.py +106 -0
  271. nbxtensor/ops/upsample_nearest2d.py +52 -0
  272. nbxtensor/ops/upsample_nearest3d.py +61 -0
  273. nbxtensor/ops/var.py +136 -0
  274. nbxtensor/ops/vector_norm.py +216 -0
  275. nbxtensor/ops/weight_norm.py +104 -0
  276. nbxtensor/ops/where.py +39 -0
  277. nbxtensor/ops/wrappers.py +825 -0
  278. nbxtensor/ops/zeros_op.py +20 -0
  279. nbxtensor/quant/__init__.py +1 -0
  280. nbxtensor/quant/quantized_tensor.py +187 -0
  281. nbxtensor/random/__init__.py +1 -0
  282. nbxtensor/random/mt_randn.py +154 -0
  283. nbxtensor/random/rng_pin.py +59 -0
  284. nbxtensor/random/rng_stream.py +118 -0
  285. nbxtensor/random/seeded_draw.py +223 -0
  286. nbxtensor/shape/__init__.py +157 -0
  287. nbxtensor/tensor.py +1867 -0
  288. nbxtensor-0.0.1.dist-info/METADATA +379 -0
  289. nbxtensor-0.0.1.dist-info/RECORD +291 -0
  290. nbxtensor-0.0.1.dist-info/WHEEL +4 -0
  291. nbxtensor-0.0.1.dist-info/licenses/LICENSE +189 -0
@@ -0,0 +1,2350 @@
1
+ """NeuroBrix kernel launcher — the dispatch layer launches, Triton compiles.
2
+
3
+ Contract: docs/internal/metal_launcher_contract.md. The owner's universal R33
4
+ (2026-09-05) counts upstream Triton's launch path (`kernel[grid]` →
5
+ `triton.runtime.driver.active`, whose CUDA backend imports torch for the
6
+ device, the stream and its benchmark buffers) as a violation on every
7
+ backend. This module keeps Triton as the COMPILER — its binder specialises
8
+ the arguments and its compiler produces the binary, both torch-free — and
9
+ hands the launch to a vendor driver of its own: the compiled binary, the
10
+ entry name, the grid, the block, the shared-memory size, the stream, and the
11
+ arguments as integer pointers and typed scalars.
12
+
13
+ launch(kernel, grid, *args, **constexprs_and_options)
14
+
15
+ is the one entry point; `install()` routes every `kernel[grid](...)` of the
16
+ house library through it (the seam for the whole-zoo bit-identity gate,
17
+ `NBX_LAUNCHER=triton` restores upstream for the differential).
18
+
19
+ CUDA is the first client (`CudaDriver`, ctypes on libcuda); ROCm and Metal
20
+ implement the same four calls behind the same `Driver` interface.
21
+ """
22
+ from __future__ import annotations
23
+
24
+ import ctypes
25
+ import os
26
+
27
+ import numpy as np
28
+ import struct
29
+ from typing import Any, Dict, List, NamedTuple, Optional, Tuple
30
+
31
+ # ---------------------------------------------------------------------------
32
+ # The target: from the engine's hardware profile, never from a driver probe
33
+ # ---------------------------------------------------------------------------
34
+
35
+ _TARGET = None # triton.backends.compiler.GPUTarget, resolved once
36
+
37
+
38
+ def _compute_capability() -> int:
39
+ """Compute capability as an int (70 for sm_70): the hardware profile the
40
+ process was given (`kernels.wrappers.set_hardware_profile`), else one
41
+ driver attribute query at first use (cached, never per launch)."""
42
+ try:
43
+ import sys as _sys
44
+ _w = _sys.modules.get("nbxtensor.ops.wrappers") # never imported from here: the
45
+ prof = getattr(_w, "_HARDWARE_PROFILE", None) # wrapper module is heavy and above us
46
+ devs = getattr(prof, "devices", None) if prof is not None else None
47
+ if devs:
48
+ cc = str(getattr(devs[0], "compute_capability", "") or "")
49
+ if cc and "." in cc:
50
+ major, minor = cc.split(".")[:2]
51
+ return int(major) * 10 + int(minor)
52
+ except Exception:
53
+ pass
54
+ return CudaDriver.instance().compute_capability()
55
+
56
+
57
+ _SMEM_LIMIT: Optional[int] = None
58
+
59
+
60
+ def max_shared_memory_per_block() -> Optional[int]:
61
+ """The executing device's shared-memory ceiling per block, or None.
62
+
63
+ The hard limit comes from the device and never from a table of names —
64
+ sm_86 declares 99 KB where sm_80 declares 163, and a table keyed on the
65
+ capability MAJOR gives the first the second's answer. That is what stopped
66
+ every language model on an A40: the flash tile was sized for sm_80 and
67
+ asked 164 352 bytes of a card that holds 101 376.
68
+
69
+ None when the driver cannot be asked (no CUDA, a backend without the
70
+ query). The caller must then leave its proposal alone rather than invent a
71
+ budget — pruning on a guessed number deletes configurations that work.
72
+
73
+ Cached for the life of the process: a driver query in a hot path is an
74
+ anti-pattern this engine has already paid for once.
75
+ """
76
+ global _SMEM_LIMIT
77
+ if _SMEM_LIMIT is None:
78
+ try:
79
+ _SMEM_LIMIT = int(active_driver().max_shared_memory_per_block())
80
+ except Exception:
81
+ return None
82
+ return _SMEM_LIMIT
83
+
84
+
85
+ def bind_target(compile_target) -> None:
86
+ """Bind this process's compile target without asking a driver: a derivation that compiles and launches nothing
87
+ (a key census) states the target its hardware description declares. `target()` answers it from then on."""
88
+ global _TARGET
89
+ _TARGET = compile_target
90
+
91
+
92
+ def target():
93
+ """The compile target of this process: vendor, capability, warp size —
94
+ from the engine's data, never from `triton.runtime.driver.active` (whose
95
+ backend probes import torch). The one source for every module that
96
+ used to ask Triton's driver (matmul / config spaces, the autotune cache).
97
+
98
+ The vendor is the DRIVER's answer, not a literal here: on Apple the
99
+ vendor is `metal`, the arch is a device name string rather than a
100
+ compute capability, and both come from the hardware profile like
101
+ everything else. `CudaDriver.target()` returns exactly what this used to
102
+ return, so the CUDA path is unchanged."""
103
+ global _TARGET
104
+ if _TARGET is None:
105
+ _TARGET = active_driver().target()
106
+ return _TARGET
107
+
108
+
109
+ def arch() -> int:
110
+ """Compute capability as an int (70 for sm_70) — `target().arch`."""
111
+ return int(target().arch)
112
+
113
+
114
+ # ---------------------------------------------------------------------------
115
+ # The vendor driver interface
116
+ # ---------------------------------------------------------------------------
117
+
118
+ def _cubin_param_count(cubin: bytes, name: str) -> Optional[int]:
119
+ """The number of kernel parameters recorded in a cubin — the count of
120
+ EIATTR_KPARAM_INFO entries in its `.nv.info.<kernel>` section.
121
+
122
+ Read from the binary because the driver on this box (535, CUDA 12.2) has
123
+ no `cuFuncGetParamInfo`, and the driver-level `launch` must still refuse
124
+ an argument list of the wrong length instead of launching it: a missing
125
+ trailing parameter is read as garbage by the kernel and faults, or worse
126
+ does not. ELF64 walk: section headers, the string table, then the
127
+ `.nv.info.<name>` attribute stream (format byte, attribute byte, a
128
+ 2-byte size for the SVAL format). Returns None when the section is
129
+ absent, and the launch then cannot check the count."""
130
+ try:
131
+ if cubin[:4] != b"\x7fELF":
132
+ return None
133
+ e_shoff = struct.unpack_from("<Q", cubin, 0x28)[0]
134
+ e_shentsize, e_shnum, e_shstrndx = struct.unpack_from("<HHH", cubin, 0x3A)
135
+
136
+ def section(i):
137
+ off = e_shoff + i * e_shentsize
138
+ sh_name, _t, _f, _a, sh_offset, sh_size = struct.unpack_from("<IIQQQQ", cubin, off)
139
+ return sh_name, sh_offset, sh_size
140
+
141
+ _, str_off, str_size = section(e_shstrndx)
142
+ strtab = cubin[str_off:str_off + str_size]
143
+ wanted = f".nv.info.{name}".encode()
144
+ for i in range(e_shnum):
145
+ sh_name, off, size = section(i)
146
+ end = strtab.index(b"\0", sh_name)
147
+ if strtab[sh_name:end] != wanted:
148
+ continue
149
+ data = cubin[off:off + size]
150
+ p, n = 0, 0
151
+ while p + 2 <= len(data):
152
+ fmt, attr = data[p], data[p + 1]
153
+ p += 2
154
+ if fmt == 0x01: # EIFMT_NVAL
155
+ pass
156
+ elif fmt in (0x02, 0x03): # EIFMT_BVAL / EIFMT_HVAL: a 16-bit value
157
+ p += 2
158
+ elif fmt == 0x04: # EIFMT_SVAL: 16-bit size, then the payload
159
+ (sz,) = struct.unpack_from("<H", data, p)
160
+ p += 2 + sz
161
+ else:
162
+ return None
163
+ if attr == 0x17: # EIATTR_KPARAM_INFO
164
+ n += 1
165
+ return n
166
+ return None
167
+ except Exception:
168
+ return None
169
+
170
+
171
+ class Driver:
172
+ """What a backend implements: load a binary, launch it, and name the
173
+ compile target it wants.
174
+
175
+ Three class attributes carry the facts the launcher used to assume were
176
+ CUDA's. Each was a CUDA literal in the launch path until the Metal port
177
+ pressed on it, and each default is the CUDA answer, so a driver that says
178
+ nothing behaves exactly as before.
179
+ """
180
+
181
+ #: The key of the compiled artifact in `CompiledKernel.asm`. CUDA emits
182
+ #: "cubin"; Metal emits "msl"; ROCm emits "hsaco".
183
+ artifact_kind = "cubin"
184
+
185
+ #: Triton's ABI (>= 3.6) passes a global-scratch and a profile-scratch
186
+ #: pointer after the kernel's own arguments. A backend whose launch ABI
187
+ #: has no such slots says so rather than receiving two stray zeros.
188
+ wants_scratch_params = True
189
+
190
+ #: A driver that binds a pointer by WRAPPING memory (Metal: a no-copy buffer per launch) needs to know how far
191
+ #: the tensor behind it reaches; it says so and receives `extents=` — the byte extent each pointer's tensor can
192
+ #: address, aligned with `params`, None where it is not a tensor. A driver that binds a raw address (CUDA) does
193
+ #: not ask and its launch is unchanged.
194
+ wants_extents = False
195
+
196
+ def load(self, binary: bytes, name: str, shared: int): # pragma: no cover - interface
197
+ raise NotImplementedError
198
+
199
+ def launch(self, function, grid, block, shared: int, stream: int, params,
200
+ names=None, types=None, trailing=None) -> None: # pragma: no cover
201
+ raise NotImplementedError
202
+
203
+ def trailing_buffers(self, metadata):
204
+ """Buffers the COMPILED kernel declares after the caller's arguments,
205
+ which only this backend knows how to build and read.
206
+
207
+ A backend whose emitter appends parameters of its own — a device-print
208
+ record area, a device-assert status area — declares them in the compile
209
+ metadata, and a driver that binds only the caller's arguments leaves
210
+ those slots unbound. Metal does not complain: `packArguments` binds the
211
+ tuple it is given at indices 0..n-1 and nothing checks the count, so a
212
+ `tl.device_assert` the author explicitly kept writes into a slot nobody
213
+ bound and the failure is LOST (measured 2026-09-17, triton-ext,
214
+ `probe_ki19_driver_protections.py`: a failing assert ran to completion
215
+ and raised nothing). CUDA has had the matching protection since this
216
+ launcher existed — `CudaDriver.launch` refuses when the parameter count
217
+ is not exactly what the cubin declares.
218
+
219
+ This is the seam for that: `prepare` asks the driver ONCE per
220
+ compilation what the kernel declares beyond the call, and `launch`
221
+ hands the answer back. The value is opaque to the launcher — naming
222
+ what is in it would put a backend's ABI in a file that must not know
223
+ one. CUDA declares nothing and returns None.
224
+ """
225
+ return None
226
+
227
+ def block_for(self, metadata):
228
+ """Threads per block, from the compiled metadata.
229
+
230
+ `num_warps * warp_size` on CUDA. Not universal: triton-msl documents
231
+ that its C++ path can overwrite `metadata.block_size` with a value
232
+ meant for a different launch shape, so the Metal driver reads the
233
+ emitted kernel's own size instead of computing one.
234
+ """
235
+ return (32 * int(metadata.num_warps), 1, 1)
236
+
237
+ def target(self): # pragma: no cover - interface
238
+ raise NotImplementedError
239
+
240
+ def max_shared_memory_per_block(self): # pragma: no cover - interface
241
+ """The device's per-block shared-memory ceiling, in bytes.
242
+
243
+ A backend that cannot ask raises, and the module-level helper turns
244
+ that into None — which every caller reads as "do not prune"."""
245
+ raise NotImplementedError
246
+
247
+
248
+ class CudaDriver(Driver):
249
+ """libcuda through ctypes. The context is the primary context of the
250
+ allocator's current device, already alive through the runtime API."""
251
+
252
+ _inst: Optional["CudaDriver"] = None
253
+ CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES = 8
254
+ CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR = 75
255
+ CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR = 76
256
+ # The OPT-IN maximum, not the 48 KB default: a kernel that asks for more
257
+ # than 48 KB must opt in per function, and this is the ceiling on that ask.
258
+ # It is the number an over-large tile is refused against.
259
+ CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN = 97
260
+
261
+ @classmethod
262
+ def instance(cls) -> "CudaDriver":
263
+ if cls._inst is None:
264
+ cls._inst = cls()
265
+ return cls._inst
266
+
267
+ def target(self):
268
+ """`GPUTarget("cuda", <capability>, 32)` — what `target()` returned
269
+ as a literal before the driver registry existed."""
270
+ from triton.backends.compiler import GPUTarget
271
+ return GPUTarget("cuda", _compute_capability(), 32)
272
+
273
+ def __init__(self) -> None:
274
+ lib = None
275
+ for name in ("libcuda.so.1", "libcuda.so"):
276
+ try:
277
+ lib = ctypes.CDLL(name)
278
+ break
279
+ except OSError:
280
+ continue
281
+ if lib is None:
282
+ raise RuntimeError("NeuroBrix launcher: libcuda not found")
283
+ self.lib = lib
284
+ self._check(lib.cuInit(0), "cuInit")
285
+ # ctypes infers the type of every argument at every call unless the signature is
286
+ # declared. The launch is called six hundred times per decoded token, so its signature
287
+ # is declared once here and the call passes plain integers (2026-09-08).
288
+ lib.cuLaunchKernel.restype = ctypes.c_int
289
+ lib.cuLaunchKernel.argtypes = [ctypes.c_void_p, ctypes.c_uint, ctypes.c_uint, ctypes.c_uint,
290
+ ctypes.c_uint, ctypes.c_uint, ctypes.c_uint, ctypes.c_uint,
291
+ ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p]
292
+ self._modules: Dict[Tuple[int, bytes], ctypes.c_void_p] = {}
293
+ self._param_counts: Dict[int, Optional[int]] = {} # function handle -> parameters the cubin declares
294
+ # (function handle, argument kinds) -> the C cells and the pointer array reused at every
295
+ # launch of that function: the buffer is written in place, never rebuilt (2026-09-08).
296
+ self._arg_buffers: Dict[Tuple[int, tuple], tuple] = {}
297
+
298
+ def _check(self, ret: int, what: str) -> None:
299
+ if ret != 0:
300
+ msg = ctypes.c_char_p()
301
+ self.lib.cuGetErrorString(ret, ctypes.byref(msg))
302
+ raise RuntimeError(f"NeuroBrix launcher: {what} failed ({ret}: "
303
+ f"{(msg.value or b'?').decode()})")
304
+
305
+ def _ensure_context(self) -> None:
306
+ ctx = ctypes.c_void_p()
307
+ self._check(self.lib.cuCtxGetCurrent(ctypes.byref(ctx)), "cuCtxGetCurrent")
308
+ if ctx.value:
309
+ return
310
+ from nbxtensor.memory.allocator import DeviceAllocator
311
+ dev = ctypes.c_int(int(DeviceAllocator.get_device()))
312
+ self._check(self.lib.cuDevicePrimaryCtxRetain(ctypes.byref(ctx), dev), "cuDevicePrimaryCtxRetain")
313
+ self._check(self.lib.cuCtxSetCurrent(ctx), "cuCtxSetCurrent")
314
+
315
+ def compute_capability(self) -> int:
316
+ from nbxtensor.memory.allocator import DeviceAllocator
317
+ dev = ctypes.c_int(int(DeviceAllocator.get_device()))
318
+ major, minor = ctypes.c_int(), ctypes.c_int()
319
+ self._check(self.lib.cuDeviceGetAttribute(ctypes.byref(major), self.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, dev), "cuDeviceGetAttribute")
320
+ self._check(self.lib.cuDeviceGetAttribute(ctypes.byref(minor), self.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, dev), "cuDeviceGetAttribute")
321
+ return major.value * 10 + minor.value
322
+
323
+ def max_shared_memory_per_block(self) -> int:
324
+ """The device's opt-in shared-memory ceiling per block, in bytes.
325
+
326
+ Asked of the DRIVER, not of a table of names. A table is right until
327
+ the next card ships; a driver query is right forever. The vendor YAMLs
328
+ keep their role — good defaults and the name in a certificate path —
329
+ but a hard limit is never one of their answers.
330
+
331
+ Read once per process and cached by the caller, never per launch.
332
+ """
333
+ from nbxtensor.memory.allocator import DeviceAllocator
334
+ dev = ctypes.c_int(int(DeviceAllocator.get_device()))
335
+ out = ctypes.c_int()
336
+ self._check(self.lib.cuDeviceGetAttribute(
337
+ ctypes.byref(out),
338
+ self.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, dev),
339
+ "cuDeviceGetAttribute")
340
+ return int(out.value)
341
+
342
+ def load(self, binary: bytes, name: str, shared: int):
343
+ self._ensure_context()
344
+ from nbxtensor.memory.allocator import DeviceAllocator
345
+ dev = int(DeviceAllocator.get_device())
346
+ key = (dev, binary)
347
+ module = self._modules.get(key)
348
+ if module is None:
349
+ module = ctypes.c_void_p()
350
+ buf = ctypes.create_string_buffer(binary, len(binary))
351
+ self._check(self.lib.cuModuleLoadData(ctypes.byref(module), buf), "cuModuleLoadData")
352
+ self._modules[key] = module
353
+ function = ctypes.c_void_p()
354
+ self._check(self.lib.cuModuleGetFunction(ctypes.byref(function), module, name.encode()), f"cuModuleGetFunction({name})")
355
+ self._param_counts[int(function.value)] = _cubin_param_count(binary, name)
356
+ if shared > 48 * 1024:
357
+ # A compiled kernel asking more shared memory than the device grants a block is a config this device
358
+ # cannot run: Triton's own CUDA launcher raises OutOfResources here, which the autotuner scores `inf`
359
+ # and passes. Asked of the driver as is, it failed `cuFuncSetAttribute (1: invalid argument)` and ended
360
+ # the sweep: the matrix-unit flash at head dim 256 on a V100 (96 KB), 2026-10-22.
361
+ limit = max_shared_memory_per_block()
362
+ if limit is not None and shared > limit:
363
+ from triton.runtime.errors import OutOfResources
364
+ raise OutOfResources(shared, limit, "shared memory")
365
+ self._check(self.lib.cuFuncSetAttribute(function, self.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, ctypes.c_int(shared)), "cuFuncSetAttribute")
366
+ return function
367
+
368
+ def launch(self, function, grid, block, shared: int, stream: int, params,
369
+ names=None, types=None, trailing=None) -> None:
370
+ # `types` is for a backend that packs its scalars into one buffer and
371
+ # needs the field offsets; CUDA binds each parameter to its own slot and
372
+ # has no use for it.
373
+ if trailing is not None:
374
+ raise RuntimeError(
375
+ "NeuroBrix launcher: trailing buffers were computed for this "
376
+ f"kernel ({trailing!r}) and the CUDA driver cannot bind them. "
377
+ "It declares none, so this is a driver/metadata mismatch, not "
378
+ "something to launch past.")
379
+ # Two refusals BEFORE anything reaches the device (the launcher
380
+ # contract's ownership rules, checked by `verify_driver_contract`):
381
+ # the argument list must be exactly what the cubin declares, and
382
+ # every non-null pointer must be an address this engine's allocator
383
+ # handed out. A kernel launched past either fault reads garbage or
384
+ # foreign memory, and a CUDA fault is sticky for the whole context.
385
+ expected = self._param_counts.get(int(function.value))
386
+ if expected is not None and len(params) != expected:
387
+ raise TypeError(f"NeuroBrix launcher: the kernel declares {expected} parameters, "
388
+ f"{len(params)} given — refused, not launched")
389
+ from nbxtensor.memory.allocator import DeviceAllocator
390
+ for i, (kind, value) in enumerate(params):
391
+ if kind == "ptr" and value and not DeviceAllocator.holds(int(value)):
392
+ # The refusal names the parameter: "which buffer" is the
393
+ # whole question the caller then has to answer (Ming's
394
+ # embedding, 2026-09-14, refused an address nobody could
395
+ # attribute from the address alone).
396
+ what = (f"parameter '{names[i]}'" if names and i < len(names)
397
+ else f"parameter #{i}")
398
+ raise ValueError(f"NeuroBrix launcher: device address {int(value):#x} ({what}) was not "
399
+ f"handed out by the allocator — refused, not launched")
400
+ # A caller may give one or two extents (the contract's checker does); CUDA wants three.
401
+ gx, gy, gz = (tuple(int(g) for g in grid) + (1, 1, 1))[:3]
402
+ bx, by, bz = (tuple(int(b) for b in block) + (1, 1, 1))[:3]
403
+ if gx * gy * gz <= 0:
404
+ return
405
+ # The argument buffer is built ONCE per (function, argument kinds) and written in place
406
+ # afterwards: a decode step launches six hundred kernels, and allocating one ctypes object
407
+ # per parameter plus a pointer array per launch was a fifth of the launcher's cost
408
+ # (2026-09-08). The buffer belongs to this driver and this function, and every launch of
409
+ # that function overwrites it before the call — a launch is synchronous from Python's side
410
+ # (the values are read by cuLaunchKernel before it returns), so one buffer is enough.
411
+ kinds = tuple(k for k, _ in params)
412
+ cache_key = (int(function.value), kinds)
413
+ slot = self._arg_buffers.get(cache_key)
414
+ if slot is None:
415
+ storage = [_ctypes_slot(kind) for kind in kinds]
416
+ arr = (ctypes.c_void_p * len(storage))(*[ctypes.addressof(c) for c in storage])
417
+ slot = self._arg_buffers[cache_key] = (storage, arr)
418
+ storage, arr = slot
419
+ for cell, (_kind, v) in zip(storage, params):
420
+ cell.value = v
421
+ rc = self.lib.cuLaunchKernel(function, gx, gy, gz, bx, by, bz, shared, stream, arr, None)
422
+ if rc != 0:
423
+ self._check(rc, "cuLaunchKernel")
424
+
425
+
426
+ # -- the contract surface (`triton/launcher_contract.py`): compile once,
427
+ # launch many times; streams and events are the allocator's. The same
428
+ # checker that gates the Metal driver gates this object, unchanged.
429
+ backend = "cuda"
430
+
431
+ def compile(self, jit_fn, signature, constexprs, num_warps: int = 4,
432
+ specialization=None, num_stages=None):
433
+ from triton.compiler import compile as triton_compile
434
+ from nbxtensor.backend.launcher_contract import ArgSlot
435
+ specialization = dict(specialization or {})
436
+ params = list(jit_fn.params)
437
+ names = [kp.name for kp in params]
438
+ unknown = set(signature) - set(names)
439
+ if unknown:
440
+ raise ValueError(f"NeuroBrix launcher: signature names {sorted(unknown)} are not parameters of {jit_fn.__name__}")
441
+ spec = []
442
+ bound = {}
443
+ for kp in params:
444
+ kind = signature.get(kp.name)
445
+ if kp.is_constexpr or kind == "constexpr":
446
+ if kp.name not in constexprs:
447
+ raise ValueError(f"NeuroBrix launcher: constexpr {kp.name!r} has no value")
448
+ bound[kp.name] = constexprs[kp.name]
449
+ spec.append(("constexpr", constexprs[kp.name]))
450
+ continue
451
+ if kind is None:
452
+ raise ValueError(f"NeuroBrix launcher: parameter {kp.name!r} is missing from the signature")
453
+ bound[kp.name] = 0
454
+ spec.append((kind, specialization.get(kp.name))) # None: no marker; "" / "D": Triton's own spellings
455
+ _, _, backend = _binder(jit_fn)
456
+ options = {"num_warps": int(num_warps)}
457
+ if num_stages is not None:
458
+ options["num_stages"] = int(num_stages)
459
+ _forward_debug(jit_fn, options)
460
+ # `_pack_args` builds the compile options from its KWARGS argument (the
461
+ # `options` one only feeds the cache key), so the same dict goes in both.
462
+ options, sig, cexprs, attrs = jit_fn._pack_args(backend, options, bound, spec, options)
463
+ src = _source(jit_fn, sig, cexprs, attrs)
464
+ # A backend that cannot emit for the device may WARN and compute on the
465
+ # CPU instead. Numbers come back, nothing raises, and the caller has an
466
+ # answer from hardware it did not ask for. Correct-or-refuse applies to
467
+ # the compile as much as to the launch.
468
+ import warnings as _warnings
469
+ with _warnings.catch_warnings(record=True) as _w:
470
+ _warnings.simplefilter("always")
471
+ compiled = triton_compile(src, target=target(), options=options.__dict__)
472
+ try:
473
+ from nbxtensor.backend.metal_backend import backend_fallback_markers
474
+ _markers = backend_fallback_markers()
475
+ except Exception: # noqa: BLE001
476
+ _markers = ()
477
+ if _markers:
478
+ for _warn in _w:
479
+ _msg = str(_warn.message)
480
+ if any(_m in _msg for _m in _markers):
481
+ raise RuntimeError(
482
+ f"the Metal backend could not emit {kernel.__name__} for "
483
+ f"this device and fell back to the CPU. Refusing the "
484
+ f"result: it would be an answer from hardware the caller "
485
+ f"did not ask for, and it is indistinguishable from a "
486
+ f"correct one. The backend said: {_msg.strip()[:400]}")
487
+ for _warn in _w: # keep every other warning visible
488
+ _warnings.warn_explicit(_warn.message, _warn.category,
489
+ _warn.filename, _warn.lineno)
490
+ md = compiled.metadata
491
+ if getattr(md, "num_ctas", 1) != 1 or getattr(md, "global_scratch_size", 0) or getattr(md, "profile_scratch_size", 0):
492
+ raise RuntimeError(f"NeuroBrix launcher: {jit_fn.__name__} needs clusters or scratch memory the CUDA client does not provide yet")
493
+ cubin = compiled.asm["cubin"]
494
+ function = CudaDriver.instance().load(cubin, md.name, md.shared)
495
+ binding = []
496
+ ordered = {}
497
+ for kp in params:
498
+ kind = signature.get(kp.name)
499
+ if kp.is_constexpr or kind == "constexpr":
500
+ continue
501
+ ordered[kp.name] = kind
502
+ binding.append(ArgSlot(index=len(binding), name=kp.name, is_pointer=kind.startswith("*"),
503
+ dtype=kind if not kind.startswith("*") else kind[1:]))
504
+ return CudaCompiledKernel(md.name, bytes(cubin), 32 * int(md.num_warps), int(md.shared), constexprs,
505
+ specialization, binding, ordered, function)
506
+
507
+ @staticmethod
508
+ def _allocator():
509
+ from nbxtensor.memory.allocator import DeviceAllocator
510
+ return DeviceAllocator
511
+
512
+ def create_stream(self) -> int:
513
+ return self._allocator().create_stream()
514
+
515
+ def destroy_stream(self, stream: int) -> None:
516
+ self._allocator().destroy_stream(stream)
517
+
518
+ def synchronize_stream(self, stream: int) -> None:
519
+ self._allocator().stream_synchronize(stream)
520
+
521
+ def create_event(self, timing: bool = False) -> int:
522
+ return self._allocator().create_event(timing=timing)
523
+
524
+ def destroy_event(self, event: int) -> None:
525
+ self._allocator().destroy_event(event)
526
+
527
+ def record_event(self, event: int, stream: int = 0) -> None:
528
+ self._allocator().record_event(event, stream)
529
+
530
+ def synchronize_event(self, event: int) -> None:
531
+ self._allocator().event_synchronize(event)
532
+
533
+ def wait_event(self, stream: int, event: int) -> None:
534
+ self._allocator().stream_wait_event(stream, event)
535
+
536
+ def elapsed_ms(self, start: int, end: int) -> float:
537
+ return self._allocator().event_elapsed_ms(start, end)
538
+
539
+ # ---------------------------------------------------------------------------
540
+ # Which driver this process launches through — the one place a backend is named
541
+ # ---------------------------------------------------------------------------
542
+ #
543
+ # `launch()` used to end in `CudaDriver.instance()`. The `Driver` base class
544
+ # was already there, so the second backend was anticipated; there was just no
545
+ # way to install one. This is that way, and it is a table rather than a
546
+ # branch: adding ROCm is a row.
547
+ #
548
+ # The backend NAME is not decided here either. It comes from the seam that
549
+ # already resolves it for the allocator, so the engine has one detection
550
+ # rather than two opinions.
551
+
552
+ _DRIVER: Optional[Driver] = None
553
+
554
+
555
+
556
+ def register_driver(driver: Optional[Driver]) -> None:
557
+ """Install the driver this process launches through. `None` clears it."""
558
+ global _DRIVER, _TARGET
559
+ _DRIVER = driver
560
+ _TARGET = None # the target is the driver's answer
561
+
562
+
563
+ def active_driver() -> Driver:
564
+ """The driver in force, resolved once from the detected backend."""
565
+ global _DRIVER
566
+ if _DRIVER is None:
567
+ _DRIVER = _resolve_driver()
568
+ return _DRIVER
569
+
570
+
571
+ def _resolve_driver() -> Driver:
572
+ from nbxtensor.backend.runtime import gpu_backend
573
+ name = gpu_backend() # a backend it cannot read raises its own error: never answered with CUDA
574
+ if name == "metal":
575
+ # WHICH Metal backend is a profile selection, and a driver implements ONE
576
+ # backend's launch ABI (ours reads the fork's MSL conventions for scalar
577
+ # binding). So the SEAM answers which module launches, or refuses by
578
+ # name. A static {"metal": <the fork's driver>} table was the previous
579
+ # answer, and it is what launched triton-ext-compiled kernels through the
580
+ # fork's ABI: scalars after the first arrived as 0, every mask went
581
+ # false, and the kernel returned exact zeros without failing.
582
+ from nbxtensor.backend.metal_backend import nbx_driver_module
583
+ from importlib import import_module
584
+ return import_module(nbx_driver_module()).driver()
585
+ if name == "cuda":
586
+ return CudaDriver.instance()
587
+ # Any other backend has no launch driver here yet: refused by name, never handed the CUDA one (an AMD card failed
588
+ # inside libcuda at its first launch, 2026-10-30). A new backend implements the driver interface (BACKENDS.md).
589
+ raise RuntimeError(f"nbxtensor launcher: no launch driver for the {name!r} backend — the library launches on cuda "
590
+ f"and metal; a backend implements the launcher's Driver interface (docs/nbxtensor/BACKENDS.md)")
591
+
592
+
593
+ def _unsupported(kind):
594
+ raise RuntimeError(f"NeuroBrix launcher: unsupported scalar kind {kind!r}")
595
+
596
+
597
+ # ---------------------------------------------------------------------------
598
+ # Argument packing: the signature dictates the kind of every parameter
599
+ # ---------------------------------------------------------------------------
600
+
601
+ _INT_KINDS = {"i1": "i8", "i8": "i8", "i16": "i16", "i32": "i32", "i64": "i64",
602
+ "u1": "u8", "u8": "u8", "u16": "u16", "u32": "u32", "u64": "u64"}
603
+
604
+
605
+ _CTYPES_SLOT = {"ptr": ctypes.c_uint64, "u64": ctypes.c_uint64, "bits64": ctypes.c_uint64,
606
+ "i64": ctypes.c_int64, "u32": ctypes.c_uint32, "bits32": ctypes.c_uint32,
607
+ "i32": ctypes.c_int32, "u16": ctypes.c_uint16, "bits16": ctypes.c_uint16,
608
+ "i16": ctypes.c_int16, "u8": ctypes.c_uint8, "i8": ctypes.c_int8,
609
+ "f64": ctypes.c_double}
610
+
611
+
612
+ def _ctypes_slot(kind: str):
613
+ """The C cell one launch parameter of this kind lives in — the same widths the packing
614
+ produced when it built a fresh object per launch."""
615
+ cell = _CTYPES_SLOT.get(kind)
616
+ if cell is None:
617
+ _unsupported(kind)
618
+ return cell()
619
+
620
+
621
+ def _pack_param(ty: str, value: Any) -> Tuple[str, Any]:
622
+ """One launch parameter as (kind, integer-or-float) — a pointer as its
623
+ address, a float scalar as the bit pattern of its storage type (what
624
+ Triton's C launcher does with pack_fp16/pack_fp32/pack_fp64)."""
625
+ if ty.startswith("*"):
626
+ return "ptr", int(value.data_ptr()) if hasattr(value, "data_ptr") else int(value)
627
+ if ty in _INT_KINDS:
628
+ return _INT_KINDS[ty], int(value)
629
+ if ty in ("fp32", "f32"):
630
+ return "bits32", struct.unpack("<I", struct.pack("<f", float(value)))[0]
631
+ if ty == "fp64":
632
+ return "bits64", struct.unpack("<Q", struct.pack("<d", float(value)))[0]
633
+ if ty == "fp16":
634
+ import numpy as np
635
+ return "bits16", int(np.array(float(value), dtype=np.float16).view(np.uint16))
636
+ if ty == "bf16":
637
+ from nbxtensor.dtype.convert import f32_to_bf16_bits
638
+ import numpy as np
639
+ return "bits16", int(f32_to_bf16_bits(np.array([float(value)], dtype=np.float32))[0])
640
+ raise RuntimeError(f"NeuroBrix launcher: cannot pack a parameter of type {ty!r}")
641
+
642
+
643
+ # ---------------------------------------------------------------------------
644
+ # The launch: Triton's binder + compiler, our driver
645
+ # ---------------------------------------------------------------------------
646
+
647
+ class _Prepared:
648
+ __slots__ = ("function", "signature", "shared", "num_warps", "block", "name",
649
+ "trailing")
650
+
651
+ def __init__(self, function, signature, shared, num_warps, name, block=None,
652
+ trailing=None):
653
+ self.function = function
654
+ self.signature = signature
655
+ self.shared = shared
656
+ self.num_warps = num_warps
657
+ self.block = block if block is not None else (32 * num_warps, 1, 1)
658
+ self.name = name
659
+ #: What the compiled kernel declares AFTER the caller's arguments, as
660
+ #: only its own driver can describe it (`Driver.trailing_buffers`).
661
+ #: Computed once per compilation, because it is a property of the
662
+ #: compilation and reading it per launch would cost a metadata walk on
663
+ #: the hot path.
664
+ self.trailing = trailing
665
+
666
+
667
+ # ---------------------------------------------------------------------------
668
+ # Specialisation: Triton's binder rules, written once in Python and measured
669
+ # against the C++ specialiser (tests/unit/kernels/test_launcher.py) — the C++
670
+ # one imports torch the moment it meets a tensor-like argument (the addendum's
671
+ # hypothesis, confirmed by the import stack on 2026-09-05).
672
+ # ---------------------------------------------------------------------------
673
+
674
+ # Triton's short element names for a pointer's dtype (torch / tl spellings in).
675
+ _POINTEE = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32", "float64": "fp64",
676
+ "int8": "i8", "int16": "i16", "int32": "i32", "int64": "i64",
677
+ "uint8": "u8", "uint16": "u16", "uint32": "u32", "uint64": "u64", "bool": "i1",
678
+ "float8_e4m3fn": "fp8e4nv", "float8_e5m2": "fp8e5", "float8_e4m3fnuz": "fp8e4b8", "float8_e5m2fnuz": "fp8e5b16",
679
+ "fp16": "fp16", "bf16": "bf16", "fp32": "fp32", "fp64": "fp64"}
680
+ _INT32_MIN, _INT32_MAX = -(2 ** 31), 2 ** 31
681
+ _INT64_MIN, _INT64_MAX = -(2 ** 63), 2 ** 63
682
+
683
+
684
+ _PTR_TY: Dict[Any, str] = {} # a tensor dtype -> the pointee spelling Triton's binder gives it
685
+
686
+
687
+ def specialize_arg(arg, specialize: bool = True, align: bool = True):
688
+ """(type, attr) of one runtime argument, exactly as Triton's binder:
689
+ a pointer is `*<dtype>` with 'D' when 16-byte aligned; an int is i32 /
690
+ i64 / u64 by range, 'D' when divisible by 16, the constant 1 when
691
+ specialised; a float is fp32; a bool is u1; None is a constexpr. The attr
692
+ is None when the parameter is not specialised at all."""
693
+ if arg is None:
694
+ return ("constexpr", None)
695
+ if hasattr(arg, "data_ptr") and hasattr(arg, "dtype"):
696
+ dt = arg.dtype
697
+ # The pointee's spelling is a property of the dtype, not of the call: computing it per
698
+ # launch cost six string operations per pointer argument, and a decode step launches six
699
+ # hundred kernels of fifteen to twenty arguments each (2026-09-08).
700
+ ty = _PTR_TY.get(dt)
701
+ if ty is None:
702
+ name = getattr(dt, "name", None) or str(dt).split(".")[-1]
703
+ ty = _PTR_TY[dt] = "*" + _POINTEE.get(str(name), str(name))
704
+ if not specialize:
705
+ return (ty, None)
706
+ return (ty, "D" if (align and int(arg.data_ptr()) % 16 == 0) else "")
707
+ if isinstance(arg, bool):
708
+ return ("u1", None)
709
+ if isinstance(arg, int):
710
+ if specialize and arg == 1:
711
+ return ("constexpr", 1)
712
+ if _INT32_MIN <= arg < _INT32_MAX:
713
+ ty = "i32"
714
+ elif _INT64_MIN <= arg < _INT64_MAX:
715
+ ty = "i64"
716
+ elif 0 <= arg < 2 ** 64:
717
+ ty = "u64"
718
+ else:
719
+ raise RuntimeError(f"NeuroBrix launcher: integer argument out of range: {arg}")
720
+ if not specialize:
721
+ return (ty, None)
722
+ return (ty, "D" if (align and arg % 16 == 0) else "")
723
+ if isinstance(arg, float):
724
+ return ("fp32", None)
725
+ try:
726
+ import triton.language as tl
727
+ if isinstance(arg, tl.dtype):
728
+ return ("constexpr", arg)
729
+ except Exception:
730
+ pass
731
+ raise TypeError(f"NeuroBrix launcher: cannot specialise an argument of type {type(arg).__name__}")
732
+
733
+
734
+ def _param_table(kernel):
735
+ """The kernel's parameter list as plain tuples, read once: a launch reads five attributes per
736
+ parameter, and there are fifteen to twenty of them on the kernels a decode step calls."""
737
+ table = _PARAM_TABLES.get(id(kernel))
738
+ if table is None:
739
+ table = _PARAM_TABLES[id(kernel)] = tuple(
740
+ (kp.name, kp.is_constexpr, kp.has_default, kp.default,
741
+ not kp.do_not_specialize, not kp.do_not_specialize_on_alignment,
742
+ getattr(kp, "annotation_type", None))
743
+ for kp in kernel.params)
744
+ return table
745
+
746
+
747
+ _PARAM_TABLES: Dict[int, tuple] = {}
748
+
749
+
750
+ def nbx_binder(kernel, args, kwargs):
751
+ """(bound_args, specialization, options) — the same triple Triton's
752
+ generated binder returns, from the kernel's parameter list."""
753
+ bound = {}
754
+ spec = []
755
+ positional = list(args)
756
+ options = dict(kwargs)
757
+ for i, (name, is_constexpr, has_default, default, specialize_p, align_p, ann) in enumerate(_param_table(kernel)):
758
+ if i < len(positional):
759
+ value = positional[i]
760
+ elif name in options:
761
+ value = options.pop(name)
762
+ elif has_default:
763
+ value = default
764
+ else:
765
+ raise TypeError(f"NeuroBrix launcher: {kernel.__name__}() missing argument {name!r}")
766
+ bound[name] = value
767
+ if is_constexpr:
768
+ spec.append(("constexpr", value))
769
+ continue
770
+ specialize = specialize_p
771
+ align = align_p
772
+ if ann:
773
+ if isinstance(ann, str) and (ann == "u1" or ann[:2] in ("fp", "bf")):
774
+ specialize = False
775
+ if specialize:
776
+ spec.append((ann,) + tuple(specialize_arg(value, True, align)[1:]))
777
+ else:
778
+ spec.append((ann, None))
779
+ continue
780
+ spec.append(specialize_arg(value, specialize, align))
781
+ if len(positional) > len(_param_table(kernel)):
782
+ raise TypeError(f"NeuroBrix launcher: {kernel.__name__}() takes {len(kernel.params)} arguments, {len(positional)} given")
783
+ return bound, spec, options
784
+
785
+
786
+ _binders: Dict[int, tuple] = {} # id(kernel) → (kernel_cache, key_cache, backend)
787
+
788
+
789
+ def reset_caches() -> None:
790
+ """Drop every compiled kernel this launcher is holding.
791
+
792
+ A `_Prepared` holds a handle the DRIVER produced — a CUDA module, a Metal
793
+ pipeline — which is bound to the device and the runtime that produced it.
794
+ Anything that replaces the runtime underneath must drop these too, or the
795
+ next launch drives a handle from a runtime that no longer exists and
796
+ refuses a perfectly good tensor.
797
+
798
+ Called by `metal_device.reset_runtime_for_tests`, which is the only thing
799
+ that replaces a runtime today. The caches and the runtime have one
800
+ lifetime; clearing a subset is worse than clearing none.
801
+ """
802
+ global _TARGET
803
+ _binders.clear()
804
+ _PARAM_TABLES.clear()
805
+ _PTR_TY.clear()
806
+ _TARGET = None
807
+
808
+
809
+ def _source(kernel, signature, constexprs, attrs):
810
+ """The kernel's AST source in its own language: Triton's `create_binder`
811
+ picks `GluonASTSource` for a `@gluon.jit` kernel (the module carries
812
+ `ttg.num-warps` and its layouts before they are verified), and this
813
+ launcher, which binds without `create_binder`, must make the same choice."""
814
+ if kernel.is_gluon():
815
+ from triton.experimental.gluon._runtime import GluonASTSource
816
+ return GluonASTSource(kernel, signature, constexprs, attrs)
817
+ from triton.compiler import ASTSource
818
+ return ASTSource(kernel, signature, constexprs, attrs)
819
+
820
+
821
+ def _binder(kernel):
822
+ b = _binders.get(id(kernel))
823
+ if b is None:
824
+ from triton.compiler import make_backend
825
+ b = _binders[id(kernel)] = ({}, {}, make_backend(target()))
826
+ return b
827
+
828
+
829
+ def _forward_debug(kernel, options: Dict[str, Any]) -> None:
830
+ """Carry the kernel's `debug` flag into the compile options, as Triton's
831
+ own `JITFunction.run` does (`kwargs.get("debug", self.debug) or
832
+ knobs.runtime.debug`). `@triton.jit(debug=True)` is what keeps a
833
+ `tl.device_assert` in the binary — the gather/scatter kernels rely on it
834
+ to trap an out-of-range index — so a launcher that dropped it would
835
+ compile the traps out silently."""
836
+ if "debug" not in options:
837
+ global _DEBUG_KNOB
838
+ if _DEBUG_KNOB is None:
839
+ _DEBUG_KNOB = False
840
+ try:
841
+ from triton import knobs
842
+ _DEBUG_KNOB = bool(knobs.runtime.debug)
843
+ except Exception:
844
+ pass
845
+ options["debug"] = bool(getattr(kernel, "debug", False)) or _DEBUG_KNOB
846
+
847
+
848
+ _DEBUG_KNOB = None # the process's Triton debug knob, read once (a launch may not import)
849
+
850
+
851
+ def prepare(kernel, args, kwargs) -> Tuple[_Prepared, Dict[str, Any]]:
852
+ """Specialise (our binder), compile (Triton's compiler with the engine's
853
+ target), load (our driver) — once per specialisation and device.
854
+
855
+ A decode step launches six hundred kernels and every one of them comes back
856
+ here: what a repeat launch pays must be a dict lookup, nothing more. The
857
+ specialisation IS the identity of the compilation, so it keys the cache
858
+ directly; Triton's `compute_cache_key` (which builds and hashes a string)
859
+ runs on a miss only, where a compile is about to happen anyway. The device
860
+ comes from the allocator's own cache — `get_device` asks the runtime, and a
861
+ runtime round trip per launch is what a launcher must not do (2026-09-08:
862
+ the launcher held ten times the decode rate of the engine).
863
+ """
864
+ from triton.compiler import compile as triton_compile
865
+ from triton.runtime.jit import compute_cache_key
866
+ kernel_cache, key_cache, backend = _binder(kernel)
867
+ kwargs = dict(kwargs)
868
+ _forward_debug(kernel, kwargs) # into kwargs: `_pack_args` parses the compile options from THEM
869
+ bound_args, specialization, options = nbx_binder(kernel, args, kwargs)
870
+ from nbxtensor.memory.allocator import DeviceAllocator
871
+ from nbxtensor.backend.runtime import _cached_device_idx
872
+ dev = _cached_device_idx()
873
+ if dev is None:
874
+ dev = int(DeviceAllocator.get_device())
875
+ fast_key = (tuple(specialization), tuple(options.items()) if options else (), dev)
876
+ prep = kernel_cache.get(fast_key)
877
+ if prep is None:
878
+ key = (compute_cache_key(key_cache, specialization, options), dev)
879
+ prep = kernel_cache.get(key)
880
+ if prep is None:
881
+ options, signature, constexprs, attrs = kernel._pack_args(backend, kwargs, bound_args, specialization, options)
882
+ src = _source(kernel, signature, constexprs, attrs)
883
+ compiled = triton_compile(src, target=target(), options=options.__dict__)
884
+ md = compiled.metadata
885
+ if getattr(md, "num_ctas", 1) != 1:
886
+ raise RuntimeError("NeuroBrix launcher: cluster launches (num_ctas > 1) are not implemented")
887
+ if getattr(md, "global_scratch_size", 0) or getattr(md, "profile_scratch_size", 0):
888
+ raise RuntimeError(f"NeuroBrix launcher: {kernel.__name__} asks for scratch memory "
889
+ f"the launcher does not provide yet")
890
+ drv = active_driver()
891
+ artifact = compiled.asm.get(drv.artifact_kind)
892
+ if artifact is None:
893
+ raise RuntimeError(
894
+ f"NeuroBrix launcher: {kernel.__name__} compiled to "
895
+ f"{sorted(compiled.asm)} but {drv.__class__.__name__} takes "
896
+ f"{drv.artifact_kind!r}")
897
+ function = drv.load(artifact, md.name, md.shared)
898
+ prep = _Prepared(function, signature, md.shared, md.num_warps, md.name,
899
+ drv.block_for(md), trailing=drv.trailing_buffers(md))
900
+ kernel_cache[key] = prep
901
+ kernel_cache[fast_key] = prep
902
+ return prep, bound_args
903
+
904
+
905
+ _TRACE = os.environ.get("NBX_LAUNCH_TRACE") # a file: one line per launch, "<kernel>\t<grid>"
906
+
907
+
908
+ # The decode replay (`neurobrix.triton.replay`) records one step's FINAL launches and
909
+ # replays them without the Python band above — the engine's fast decode path, worth 8x
910
+ # on a text row. It used to record by wrapping Triton's `CompiledKernel.run`; nothing on
911
+ # the Triton branch calls that any more (R33, third peel 2026-09-05), so the recorder
912
+ # needs a seam HERE or it records a step with no kernel in it. Armed for the window of
913
+ # one recorded step and cleared after: the hot path pays one global read and a None test.
914
+ _RECORDER = None
915
+
916
+
917
+ def set_launch_recorder(fn) -> None:
918
+ """Record every launch this launcher issues: `fn(prepared, grid, params)`, with the
919
+ launch as the DRIVER will take it — the resolved kernel, three extents, and the
920
+ parameters already packed. `None` stops recording."""
921
+ global _RECORDER
922
+ _RECORDER = fn
923
+
924
+
925
+ def launch(kernel, grid, *args, **kwargs):
926
+ """The one entry point: `launch(kernel, grid, *args, **constexprs_and_options)`.
927
+
928
+ A grid with no program launches nothing, as Triton's own launchers do (`gridX * gridY * gridZ > 0`): an
929
+ element-wise op on an empty tensor (numel 0 -> cdiv(0, BLOCK) = 0 programs) computes nothing, and its tensors may
930
+ hold no buffer at all — a backend whose driver refuses a null pointer (triton-ext on Metal) refused the whole run
931
+ for it (Qwen3-VL-30B-A3B-Thinking triton, an add over a (0, 2048) tensor, Mac 2026-10-23)."""
932
+ if not callable(grid) and 0 in tuple(int(g) for g in grid):
933
+ return None
934
+ prep, bound_args = prepare(kernel, args, kwargs)
935
+ if _TRACE:
936
+ with open(_TRACE, "a") as fh:
937
+ fh.write(f"{kernel.__name__}\t{grid if not callable(grid) else 'callable'}\n")
938
+ if callable(grid):
939
+ grid = grid(bound_args)
940
+ grid = tuple(int(g) for g in grid) + (1,) * (3 - len(grid))
941
+ if 0 in grid:
942
+ return None
943
+ runtime = [(name, ty) for name, ty in prep.signature.items()
944
+ if ty != "constexpr"]
945
+ params = [_pack_param(ty, bound_args[name]) for name, ty in runtime]
946
+ # The parameter NAMES, in the same order. A backend whose artifact
947
+ # declares its arguments by name — and may declare only the ones the
948
+ # compiled kernel kept — binds by name rather than by position, which is
949
+ # the only mapping that stays correct when the two lists differ in length.
950
+ names = [name for name, _ty in runtime]
951
+ drv = active_driver()
952
+ if drv.wants_scratch_params:
953
+ params.append(("ptr", 0)) # global scratch (Triton ≥ 3.6 ABI)
954
+ params.append(("ptr", 0)) # profile scratch
955
+ if _RECORDER is not None:
956
+ _RECORDER(prep, grid, params)
957
+ # The Triton TYPE of every runtime parameter, in the same order. A backend
958
+ # that packs its scalars into one buffer computes the field offsets from
959
+ # these; our `(kind, value)` pairs have already lost the distinction between
960
+ # an i16 and a bf16, and packing to the wrong offset is silent.
961
+ extra = {}
962
+ if getattr(drv, "wants_extents", False):
963
+ # One extent per runtime parameter (the scratch slots, if any, after them have none).
964
+ extents = [_addressable_extent(bound_args[name]) if str(ty).startswith("*") else None
965
+ for name, ty in runtime]
966
+ extra["extents"] = extents + [None] * (len(params) - len(extents))
967
+ drv.launch(prep.function, grid, prep.block, prep.shared, _stream(), params,
968
+ names=names, types=[ty for _name, ty in runtime],
969
+ trailing=prep.trailing, **extra)
970
+
971
+
972
+ def _addressable_extent(value) -> Optional[int]:
973
+ """The bytes a kernel can address from a tensor argument's data pointer: the last element its shape and strides
974
+ reach, plus one, times the element size. None for a value that is not a tensor (a raw address, a scalar). A
975
+ negative stride reaches below the pointer and is not described by a forward extent: None, so the driver keeps
976
+ its whole-allocation binding. A zero-element tensor addresses nothing: 0."""
977
+ if not (hasattr(value, "data_ptr") and hasattr(value, "stride") and hasattr(value, "element_size")):
978
+ return None
979
+ shape = tuple(int(d) for d in value.shape)
980
+ if any(d == 0 for d in shape):
981
+ return 0
982
+ strides = tuple(int(s) for s in value.stride())
983
+ if any(s < 0 for s in strides):
984
+ return None
985
+ last = sum((d - 1) * s for d, s in zip(shape, strides))
986
+ return (last + 1) * int(value.element_size())
987
+
988
+
989
+ def _stream() -> int:
990
+ """The legacy default stream (0): every house kernel launches on it today,
991
+ and the allocator's copies order against it."""
992
+ return 0
993
+
994
+
995
+
996
+ # ---------------------------------------------------------------------------
997
+ # The autotune correctness screen: every candidate is checked before any is timed
998
+ # ---------------------------------------------------------------------------
999
+ #
1000
+ # An autotuner ranks configs by speed. That is safe only while every config
1001
+ # computes the same thing, and on a backend where one does not, speed is
1002
+ # exactly the wrong tiebreak: a kernel that writes half its output does half
1003
+ # the stores, so it is genuinely faster, so it wins. Measured 2026-09-07 on
1004
+ # Apple — fp16 mm [64,32]@[32,64] selected a config that left 32 of 64
1005
+ # columns zero, deterministically, with no error raised anywhere.
1006
+ #
1007
+ # Nothing here names a backend, and the tolerance is read from the hardware
1008
+ # profile rather than written down: what separates "a different summation
1009
+ # order" from "a different answer" is a property of the device's arithmetic,
1010
+ # and the profile is where the engine keeps those.
1011
+ #
1012
+ # A config that diverges is excluded from the timing and recorded with its
1013
+ # shape and its deviation. It is never dropped silently, because a config
1014
+ # quietly missing from a sweep looks like a config that lost on speed.
1015
+
1016
+ class ScreenedOut(NamedTuple):
1017
+ """One config the screen refused, and why."""
1018
+ kernel: str
1019
+ key: tuple
1020
+ config: str
1021
+ dtype: str
1022
+ deviation: float
1023
+ tolerance: float
1024
+
1025
+
1026
+ class Unscreened(NamedTuple):
1027
+ """One key where a configuration was seated WITHOUT an oracle.
1028
+
1029
+ Not a failure and not a refusal: the engine must run. It is a PROVENANCE.
1030
+ "This configuration was validated" and "this configuration was the fastest
1031
+ among candidates nobody verified" are different statements, and until this
1032
+ record existed they were written the same way and read the same way.
1033
+ """
1034
+ kernel: str
1035
+ key: tuple
1036
+ candidates: int
1037
+ reason: str
1038
+
1039
+
1040
+ _SCREENED: List[ScreenedOut] = []
1041
+ _UNSCREENED: List[Unscreened] = []
1042
+ _SCREEN_CACHE: Dict[int, set] = {} # id(tuner) -> keys already screened
1043
+
1044
+
1045
+ def screened_out() -> List[ScreenedOut]:
1046
+ """Every config the correctness screen excluded, in order."""
1047
+ return list(_SCREENED)
1048
+
1049
+
1050
+ #: Keys the screen adjudicated, with what adjudicated them — the converse of
1051
+ #: `_UNSCREENED`: a record that says only what was NOT verified lets a silent
1052
+ #: entry read as verified. `capture()` stamps these `screened: True` with the
1053
+ #: adjudicator's name.
1054
+ _ADJUDICATED: dict = {}
1055
+
1056
+
1057
+ def adjudicated() -> dict:
1058
+ return dict(_ADJUDICATED)
1059
+
1060
+
1061
+ def unscreened() -> List[Unscreened]:
1062
+ """Every key whose seated configuration no oracle adjudicated.
1063
+
1064
+ Read this wherever a chosen configuration is RECORDED, so the record
1065
+ carries what it is. The certified directory is filled only by
1066
+ `neurobrix autotune certify`, which runs its own fp64 oracle over every
1067
+ candidate; nothing on this path may ever reach it.
1068
+ """
1069
+ return list(_UNSCREENED)
1070
+
1071
+
1072
+ def clear_screened() -> None:
1073
+ _SCREENED.clear()
1074
+ _UNSCREENED.clear()
1075
+ _SCREEN_CACHE.clear()
1076
+
1077
+
1078
+ def _no_oracle_reason(oracle) -> str:
1079
+ """Why no oracle adjudicated this key — the provider, or its answer."""
1080
+ if _SCREEN_ORACLE is None:
1081
+ return "no oracle provider is installed"
1082
+ if oracle is None:
1083
+ return ("the oracle provider covers no oracle for this kernel "
1084
+ "(the GEMM class and, since 2026-09-13, the convolution family "
1085
+ "are covered; anything else is decided by the bare vote)")
1086
+ try:
1087
+ from nbxtensor.ops.screen_oracle import last_refusal
1088
+ why = last_refusal()
1089
+ except Exception: # noqa: BLE001
1090
+ why = None
1091
+ return why or "the oracle produced no reference"
1092
+
1093
+
1094
+
1095
+ def _needs_windowing(buffers, budget: int) -> Optional[str]:
1096
+ """Why the FULL screen must not run on these buffers — None when it may.
1097
+
1098
+ The budget is a rule about COST, and the full screen's cost is not the buffers' device
1099
+ bytes: it snapshots every buffer to the host and computes the reference in float64 —
1100
+ every element at 8 bytes. A convolution whose OUTPUT is small and whose INPUT is not
1101
+ (128 -> 3 channels at 2048x1024: 12 MB out, 537 MB in) passed the device-byte test and
1102
+ then held 2.1 GB of float64 input on the host before the first tap; on a 24 GB Apple
1103
+ machine that was a 27.5 GB footprint and a kill (the Mac, 2026-09-25). So the full screen
1104
+ is taken only when its float64 footprint fits the budget too; otherwise the windowed
1105
+ screen, whose cost is set by the windows and not by the shape. On a 32 MiB budget this
1106
+ moves nearly every convolution whose input exceeds 4 MiB to the windowed screen."""
1107
+ from nbxtensor.ops import screen_oracle as _so_
1108
+ total = sum(nbytes for _a, nbytes, _d in buffers)
1109
+ if total > budget:
1110
+ return f"the arguments total {total} bytes, over the screening budget {budget}"
1111
+ footprint = _so_.fp64_footprint(buffers)
1112
+ if footprint > budget:
1113
+ return (f"the reference's float64 footprint is {footprint} bytes for {total} bytes of "
1114
+ f"arguments, over the screening budget {budget}")
1115
+ return None
1116
+
1117
+
1118
+ def _screen_on_windows(tuner, kernel_name, _rk, configs, args, meta, buffers, budget, total):
1119
+ """Screen an over-budget shape on bounded row windows of its output. None if impossible.
1120
+
1121
+ The candidates' own bytes are sliced to the same rows as the oracle, so both sides are the
1122
+ same rows of the same tensor. Only the OUTPUT is windowed and compared; the inputs are not
1123
+ re-read, which is the one thing this screen checks less than the full one — said in the
1124
+ seat line rather than left to be discovered."""
1125
+ from nbxtensor.ops import screen_oracle as _so
1126
+ out_name = None
1127
+ entry = _so.ORACLES.get(getattr(getattr(tuner, "base_fn", None), "__name__", "") or "")
1128
+ if entry is not None:
1129
+ out_name = entry[1]
1130
+ named = {**dict(getattr(tuner, "nargs", None) or {}), **dict(meta or {})}
1131
+ out_tensor = named.get(out_name) if out_name else None
1132
+ if out_tensor is None:
1133
+ return None
1134
+ try:
1135
+ row_cost = _so.oracle_row_bytes(tuner, meta)
1136
+ except ValueError as exc:
1137
+ return _seat_unscreened(kernel_name, _rk, configs, len(configs),
1138
+ f"over the screening budget ({total} bytes), and {exc}")
1139
+ wins = _row_windows_for(out_tensor, budget, row_cost)
1140
+ if not wins:
1141
+ return None
1142
+ flat = _flat_rows(out_tensor)
1143
+ if flat is None:
1144
+ return None
1145
+ try:
1146
+ M, N, itemsize = flat
1147
+ row_bytes = N * itemsize
1148
+ out_addr = int(out_tensor.data_ptr())
1149
+ except Exception: # noqa: BLE001
1150
+ return None
1151
+ ranges = _rows_to_bytes(wins, row_bytes)
1152
+ where = _describe_windows(wins, M)
1153
+
1154
+ shots = []
1155
+ for config in configs:
1156
+ try:
1157
+ tuner.fn.run(*args, **{**dict(meta or {}), **config.all_kwargs()})
1158
+ shots.append((config, _snapshot_ranges(out_addr, ranges)))
1159
+ except Exception: # noqa: BLE001
1160
+ pass
1161
+ if len(shots) < 2:
1162
+ return _seat_unscreened(kernel_name, _rk, configs, len(shots),
1163
+ f"over the screening budget ({total} bytes) and fewer than "
1164
+ f"two candidates ran on the windows {where}")
1165
+
1166
+ dtype_name = None
1167
+ for _a, _n, dt in buffers:
1168
+ if _a == out_addr:
1169
+ dtype_name = dt
1170
+ break
1171
+
1172
+ # ONE ORACLE PER WINDOW, never over their span. Passing (first_start, last_end) asks for
1173
+ # every row between the first and last window — the whole product, which is the 18 GB
1174
+ # computation the windowing exists to avoid, and it fails silently into the consensus
1175
+ # path. Measured 2026-09-23: matmul_kernel fell through to consensus for exactly this
1176
+ # reason while sitting in ROW_WINDOWABLE.
1177
+ kept = None
1178
+ _so.set_last_refusal(None)
1179
+ try:
1180
+ import numpy as _np
1181
+ parts = []
1182
+ for r0, r1 in wins:
1183
+ piece = _so.windowed_reference(tuner, meta, (r0, r1))
1184
+ if piece is None:
1185
+ parts = None
1186
+ break
1187
+ # CAST TO THE OUTPUT'S DTYPE before comparing, exactly as the full-screen
1188
+ # provider does. The reference is computed in float64; a candidate's bytes are
1189
+ # fp16. Comparing them raw fails with a 4x shape mismatch — which is what the
1190
+ # swallowed exception was hiding until it was made to speak.
1191
+ if dtype_name in ("bf16", "bfloat16"):
1192
+ from nbxtensor.dtype.convert import f32_to_bf16_bits
1193
+ want = _np.ascontiguousarray(
1194
+ f32_to_bf16_bits(_np.ascontiguousarray(piece, dtype=_np.float32)))
1195
+ else:
1196
+ want = _np.ascontiguousarray(piece.astype(_so._NP.get(dtype_name, _np.float32)))
1197
+ parts.append(want.tobytes())
1198
+ if parts is not None:
1199
+ ref_bytes = b"".join(parts)
1200
+ kept = [(c, blob) for c, blob in shots
1201
+ if configs_agreeing_with_oracle([(c, blob)], ref_bytes, dtype_name)]
1202
+ else:
1203
+ print(f"[AUTOTUNE_SCREEN_WINDOWED] {kernel_name}: the windowed oracle returned "
1204
+ f"nothing for a window of this key"
1205
+ + (f" ({_so.last_refusal()})" if _so.last_refusal() else "")
1206
+ + "; falling back to consensus on the same windows.", flush=True)
1207
+ except Exception as exc: # noqa: BLE001
1208
+ # SAY why. An `except: pass` here is the exact shape this session spent a day
1209
+ # removing elsewhere: it turns "the oracle could not be computed" into "no oracle
1210
+ # covers this kernel", which is a different and false statement.
1211
+ print(f"[AUTOTUNE_SCREEN_WINDOWED] {kernel_name}: the windowed oracle raised "
1212
+ f"({type(exc).__name__}: {str(exc)[:160]}); falling back to consensus on the "
1213
+ f"same windows.", flush=True)
1214
+ kept = None
1215
+ if kept is not None:
1216
+ if not kept:
1217
+ raise RuntimeError(
1218
+ f"NeuroBrix autotune screen: {kernel_name} at key {_rk} — the fp64 oracle "
1219
+ f"contradicts EVERY candidate on {where}. Screened on windows because the "
1220
+ f"arguments ({total} bytes) exceed the screening budget; a window that "
1221
+ f"disagrees is a disagreement. Refusing to seat any of them.")
1222
+ print(f"[AUTOTUNE_SCREEN_WINDOWED] {kernel_name} at key {_rk}: arguments total "
1223
+ f"{total} bytes, over the screening budget {budget}. Screened against the fp64 "
1224
+ f"oracle on {where} — {len(kept)} of {len(shots)} candidates kept. VERIFIED ON "
1225
+ f"THOSE WINDOWS ONLY; the rest of the output and the input buffers were not "
1226
+ f"compared.", flush=True)
1227
+ _record_screen_windows(_rk, where, "fp64 oracle", len(kept), len(shots))
1228
+ return [c for c, _ in kept]
1229
+
1230
+ # No windowed oracle for this kernel: cluster the candidates by agreement on the windows.
1231
+ # Strictly less than an oracle and strictly more than nothing, and it says which it is.
1232
+ groups = {}
1233
+ for c, blob in shots:
1234
+ groups.setdefault(blob, []).append(c)
1235
+ best = max(groups.values(), key=len)
1236
+ refused = _so.last_refusal()
1237
+ why = (f"the row-windowed oracle REFUSED this key ({refused})" if refused
1238
+ else "no row-windowed oracle covers this kernel")
1239
+ print(f"[AUTOTUNE_SCREEN_WINDOWED] {kernel_name} at key {_rk}: arguments total {total} "
1240
+ f"bytes, over the screening budget {budget}, and {why}. Screened by CONSENSUS on "
1241
+ f"{where} — {len(best)} of {len(shots)} agree. No oracle adjudicated this key.", flush=True)
1242
+ _record_screen_windows(_rk, where,
1243
+ f"consensus (windowed oracle refused: {refused})" if refused
1244
+ else "consensus (no windowed oracle)", len(best), len(shots))
1245
+ return best
1246
+
1247
+
1248
+ #: What each windowed screen verified, so a proof can name it.
1249
+ SCREEN_WINDOWS: dict = {}
1250
+
1251
+
1252
+ def _record_screen_windows(key, where: str, by: str, kept: int, candidates: int) -> None:
1253
+ SCREEN_WINDOWS[str(key)] = {"windows": where, "adjudicated_by": by,
1254
+ "kept": kept, "candidates": candidates}
1255
+
1256
+
1257
+ def screen_windows_of(key):
1258
+ """The windows a key was screened on, for the proof. None when it was screened whole."""
1259
+ return SCREEN_WINDOWS.get(str(key))
1260
+
1261
+
1262
+ def _seat_unscreened(kernel: str, key, configs, candidates: int, reason: str):
1263
+ """Return the configs, having said plainly that nothing verified them.
1264
+
1265
+ The bare screen may still rank by speed — the engine never refuses to run.
1266
+ What it may no longer do is produce a line that reads as a validation.
1267
+ """
1268
+ _UNSCREENED.append(Unscreened(kernel, key, candidates, reason))
1269
+ print(f"[AUTOTUNE_UNSCREENED] {kernel} at key {key}: {reason}. The "
1270
+ f"configuration seated here is the FASTEST AMONG {candidates} "
1271
+ f"CANDIDATES THAT NOTHING VERIFIED — it is not a validated setting "
1272
+ f"and it is never written to the certified directory.", flush=True)
1273
+ return configs
1274
+
1275
+
1276
+ def _screen_rtol(dtype_name: str):
1277
+ """The tolerance for one dtype, from the hardware profile. None = exact."""
1278
+ from nbxtensor.ops._configs import active_vendor_profile
1279
+
1280
+ table = active_vendor_profile().get("autotune_screen_rtol")
1281
+ if table is None:
1282
+ raise RuntimeError(
1283
+ "the hardware profile declares no `autotune_screen_rtol`: the "
1284
+ "autotune correctness screen will not invent a tolerance, and "
1285
+ "without one it cannot tell a reordered sum from a wrong answer")
1286
+ return table.get(dtype_name)
1287
+
1288
+
1289
+ def _writable_buffers(values):
1290
+ """Every device buffer among these arguments, with its byte length.
1291
+
1292
+ Returns None when any of them is a view WITH GAPS, and the caller
1293
+ then skips screening rather than guessing. A permutation of a dense
1294
+ buffer (a KRSC conv weight read as K, C, R, S since 9c2ab78e) has none:
1295
+ its span from `data_ptr()` is exactly its bytes, so it is kept.
1296
+
1297
+ The snapshot and restore below copy a CONTIGUOUS span from `data_ptr()`.
1298
+ For a strided view that span is not the tensor: it covers the gaps
1299
+ between the view's elements, which belong to other tensors in the same
1300
+ allocation. Restoring it writes stale bytes over live data somewhere
1301
+ else entirely — which is exactly what happened, and it showed up as eight
1302
+ unrelated tests failing in the full suite while passing alone.
1303
+ """
1304
+ out = []
1305
+ for value in values:
1306
+ if not (hasattr(value, "data_ptr") and hasattr(value, "_nbytes")):
1307
+ continue
1308
+ contiguous = getattr(value, "is_contiguous", None)
1309
+ nbytes = int(value._nbytes)
1310
+ if callable(contiguous) and not contiguous() and not _dense(value):
1311
+ if not _dense_broadcast(value):
1312
+ return None
1313
+ # a BROADCAST of a dense buffer (stride 0 on the repeated axes: the attention wrapper's zero / causal bias
1314
+ # read as [B, H, Sq, Sk]) is the bytes it reaches, not its element count's — the span snapshotted and
1315
+ # restored is exactly its buffer (the matrix-unit flash's certification refused it, 2026-10-21)
1316
+ nbytes = (sum((int(n) - 1) * int(st) for n, st in zip(value.shape, value.stride())) + 1) * \
1317
+ int(value.element_size())
1318
+ out.append((int(value.data_ptr()), nbytes, getattr(getattr(value, "dtype", None), "name", "?")))
1319
+ return out
1320
+
1321
+
1322
+ def _dense_broadcast(value):
1323
+ """A view that repeats a dense buffer along stride-0 axes: its other axes form a gap-free layout (`_dense`)."""
1324
+ kept = [(sh, st) for sh, st in zip(value.shape, value.stride()) if st != 0 and sh != 1]
1325
+ expected = 1
1326
+ for st, sh in sorted((st, sh) for sh, st in kept):
1327
+ if st != expected:
1328
+ return False
1329
+ expected *= sh
1330
+ return any(st == 0 and sh != 1 for sh, st in zip(value.shape, value.stride()))
1331
+
1332
+
1333
+ def _dense(value):
1334
+ """Whether a view's elements fill a span with no gap and no overlap: its
1335
+ strides, smallest first (size-1 axes aside), are 1 then the running
1336
+ product of the sizes — a permutation of a contiguous layout."""
1337
+ expected = 1
1338
+ for st, sh in sorted((st, sh) for sh, st in zip(value.shape, value.stride()) if sh != 1):
1339
+ if st != expected:
1340
+ return False
1341
+ expected *= sh
1342
+ return True
1343
+
1344
+
1345
+
1346
+ #: How many row windows an over-budget screen takes, and where. The LAST window is anchored
1347
+ #: at the final row deliberately: the failure class that motivates screening a large shape is
1348
+ #: index overflow, and it shows at the largest linear index or nowhere (the rack, 2026-09-23).
1349
+ _SCREEN_WINDOWS = 3
1350
+
1351
+
1352
+ def _screen_windows() -> int:
1353
+ """How many row windows an over-budget screen takes. `NBX_SCREEN_WINDOWS` is the door.
1354
+
1355
+ A tunable with no door cannot be tuned, and the instruction was to set this FROM a
1356
+ measurement. Read WHEN USED, never frozen at import — a value frozen from the environment
1357
+ at import is the same defect as a literal standing in for a runtime value, and this package
1358
+ has been bitten by exactly that (`autotune_cache._dir` says so in its own docstring).
1359
+
1360
+ WHAT THE COUNT DOES, AND DOES NOT, measured 2026-09-23. `per` is
1361
+ `budget // (n_win * oracle_row_bytes)`, so more windows means proportionally smaller ones
1362
+ and the oracle's TOTAL work does not move (M=100 000, N=64, 4 MiB budget):
1363
+
1364
+ windows rows each total rows places
1365
+ 1 8192 8192 1
1366
+ 3 2730 8190 3
1367
+ 8 1024 8192 8
1368
+
1369
+ **So this is not a cost knob.** It changes WHERE the screen looks, not what it costs; the
1370
+ cost is set by the screening budget. A cell asserting the opposite went red on its first
1371
+ run, which is how the property was found.
1372
+
1373
+ WHAT THE ORACLE COSTS, chatterbox on a 16 GB V100 where its 32g-certified keys do not serve
1374
+ so all 65 keys sweep:
1375
+
1376
+ oracle OFF 379.9, 378.9 s spread 1.0 s
1377
+ oracle ON 500.3, 395.5 s spread 104.8 s
1378
+
1379
+ Every ON run is slower than every OFF run — 1.04x at the closest bound, 1.32x at the widest,
1380
+ 1.18x on the means — but the ON arm's own spread is 105 s against the OFF arm's 1 s, so the
1381
+ variance lives in the oracle path and a point factor from two reps is not a number. The
1382
+ Apple side measured 31.4x for ONE selection on a 4.6 GB shape, so the factor scales with
1383
+ what is windowed and there is no constant to adopt.
1384
+ """
1385
+ import os as _os
1386
+ raw = _os.environ.get("NBX_SCREEN_WINDOWS")
1387
+ if raw is None:
1388
+ return _SCREEN_WINDOWS
1389
+ try:
1390
+ n = int(raw)
1391
+ except ValueError:
1392
+ raise ValueError(f"NBX_SCREEN_WINDOWS={raw!r} is not an integer") from None
1393
+ if n < 1:
1394
+ raise ValueError(f"NBX_SCREEN_WINDOWS={n} would screen nothing; 1 is the minimum")
1395
+ return n
1396
+
1397
+
1398
+ def _flat_rows(out_tensor):
1399
+ """A contiguous output read as `[M, N]` rows: `(M, N, itemsize)`, or None. A 2-D output is
1400
+ the matrix; an N-D one (the convolution family's `[N, Cout, Ho, Wo]`) is
1401
+ `[prod(leading), last]` — rows of the same contiguous memory, so the candidates' bytes and
1402
+ the reference slice identically. Until 2026-09-26 only 2-D outputs qualified and every conv
1403
+ key over the budget was seated without a screen."""
1404
+ try:
1405
+ shape = tuple(int(x) for x in out_tensor.shape)
1406
+ itemsize = int(out_tensor._nbytes) // max(1, int(np.prod(shape)))
1407
+ except Exception: # noqa: BLE001
1408
+ return None
1409
+ if len(shape) < 2 or itemsize <= 0:
1410
+ return None
1411
+ M, N = int(np.prod(shape[:-1])), shape[-1]
1412
+ if M <= 0 or N <= 0:
1413
+ return None
1414
+ return M, N, itemsize
1415
+
1416
+
1417
+ def _row_windows_for(out_tensor, budget_bytes: int, oracle_row_cost: Optional[int] = None):
1418
+ """Row ranges of an output whose total bytes fit the budget, or None.
1419
+
1420
+ None means the shape cannot be windowed by rows — the caller then says so rather than
1421
+ pretending it screened something. `oracle_row_cost` is what ONE row costs the reference on
1422
+ the host when the kernel knows better than "the row in float64" (the convolution family:
1423
+ its receptive field, not its output row)."""
1424
+ flat = _flat_rows(out_tensor)
1425
+ if flat is None:
1426
+ return None
1427
+ M, N, itemsize = flat
1428
+ row_bytes = N * itemsize
1429
+ if row_bytes <= 0 or M <= 0:
1430
+ return None
1431
+ # Size the window by what the ORACLE costs, not by what the output costs. The reference
1432
+ # is computed in float64 — 8 bytes a element against 2 for an fp16 output — so a window
1433
+ # sized to the budget in output bytes needs four times the budget to adjudicate, and the
1434
+ # oracle then raises and the screen silently degrades to consensus (measured 2026-09-23,
1435
+ # matmul_kernel falling through while sitting in ROW_WINDOWABLE). The reference row is
1436
+ # what it will actually occupy.
1437
+ oracle_row_bytes = max(N * 8, int(oracle_row_cost or 0))
1438
+ n_win = _screen_windows()
1439
+ per = max(1, int(budget_bytes) // (n_win * max(oracle_row_bytes, 1)))
1440
+ if per >= M:
1441
+ return [(0, M)]
1442
+ # `n_win` windows spread across the rows, the LAST always anchored at the final row:
1443
+ # the failure class that motivates screening a large shape is index overflow, and it shows
1444
+ # at the largest linear index or nowhere.
1445
+ if n_win == 1:
1446
+ wins = [(M - per, M)]
1447
+ else:
1448
+ step = (M - per) / (n_win - 1)
1449
+ wins = [(int(round(i * step)), int(round(i * step)) + per) for i in range(n_win - 1)]
1450
+ wins.append((M - per, M))
1451
+ # de-duplicate and order; overlapping windows on a short M collapse to fewer
1452
+ out, seen = [], set()
1453
+ for r0, r1 in wins:
1454
+ r0, r1 = max(0, r0), min(M, r1)
1455
+ if r1 > r0 and (r0, r1) not in seen:
1456
+ seen.add((r0, r1)); out.append((r0, r1))
1457
+ return out or None
1458
+
1459
+
1460
+ def _describe_windows(wins, M):
1461
+ return "rows " + "; ".join(f"{r0}-{r1}" for r0, r1 in wins) + f" of {M}"
1462
+
1463
+
1464
+ def _rows_to_bytes(wins, row_bytes):
1465
+ return [(r0 * row_bytes, (r1 - r0) * row_bytes) for r0, r1 in wins]
1466
+
1467
+
1468
+ def _snapshot_ranges(address, ranges):
1469
+ """Copy only the given byte ranges of one device buffer to the host."""
1470
+ import ctypes
1471
+
1472
+ from nbxtensor.memory.allocator import DeviceAllocator
1473
+ blobs = []
1474
+ for off, length in ranges:
1475
+ host = (ctypes.c_char * length)()
1476
+ DeviceAllocator.memcpy(ctypes.addressof(host), address + off, length, kind=2)
1477
+ blobs.append(bytes(host))
1478
+ return b"".join(blobs)
1479
+
1480
+
1481
+ def _snapshot(buffers):
1482
+ import ctypes
1483
+
1484
+ from nbxtensor.memory.allocator import DeviceAllocator
1485
+
1486
+ shots = []
1487
+ for address, nbytes, _name in buffers:
1488
+ host = (ctypes.c_char * nbytes)()
1489
+ DeviceAllocator.memcpy(ctypes.addressof(host), address, nbytes, kind=2)
1490
+ shots.append(bytes(host))
1491
+ return shots
1492
+
1493
+
1494
+ def _restore(buffers, shots):
1495
+ import ctypes
1496
+
1497
+ from nbxtensor.memory.allocator import DeviceAllocator
1498
+
1499
+ for (address, nbytes, _name), blob in zip(buffers, shots):
1500
+ host = (ctypes.c_char * nbytes).from_buffer_copy(blob)
1501
+ DeviceAllocator.memcpy(address, ctypes.addressof(host), nbytes, kind=1)
1502
+
1503
+
1504
+ #: The engine spells a dtype "fp16"; the hardware profiles and numpy spell it
1505
+ #: "float16". One table, so the screen cannot silently fail to find a
1506
+ #: tolerance and fall back to comparing floats bit-for-bit — which it did on
1507
+ #: first run, excluding seven correct configs.
1508
+ _DTYPE_CANON = {
1509
+ "fp16": "float16", "float16": "float16", "half": "float16",
1510
+ "bf16": "bfloat16", "bfloat16": "bfloat16",
1511
+ "fp32": "float32", "float32": "float32", "f32": "float32",
1512
+ "fp64": "float64", "float64": "float64", "f64": "float64",
1513
+ "i8": "int8", "int8": "int8", "i16": "int16", "int16": "int16",
1514
+ "i32": "int32", "int32": "int32", "i64": "int64", "int64": "int64",
1515
+ "u8": "uint8", "uint8": "uint8", "i1": "bool_", "bool": "bool_",
1516
+ "bool_": "bool_",
1517
+ }
1518
+
1519
+
1520
+ def _as_float64(blob: bytes, canon: str):
1521
+ """One buffer's contents as float64, whatever the engine calls its type.
1522
+
1523
+ bfloat16 is not a numpy dtype: it is read as its 16 bits and widened by
1524
+ placing them in the high half of a float32, which is exactly what the
1525
+ format is.
1526
+ """
1527
+ import numpy as np
1528
+
1529
+ if canon == "bfloat16":
1530
+ bits = np.frombuffer(blob, dtype=np.uint16).astype(np.uint32) << 16
1531
+ return bits.view(np.float32).astype(np.float64)
1532
+ return np.frombuffer(blob, dtype=np.dtype(canon)).astype(np.float64)
1533
+
1534
+
1535
+ def _deviation(a: bytes, b: bytes, dtype_name: str):
1536
+ """(deviation, tolerance) between two results of the same buffer.
1537
+
1538
+ Integers and booleans are compared bit-identically: no valid reordering
1539
+ changes an integer. Floats are compared relative to the reference's own
1540
+ magnitude, against the profile's tolerance.
1541
+ """
1542
+ import numpy as np
1543
+
1544
+ canon = _DTYPE_CANON.get(str(dtype_name).lower())
1545
+ if canon is None:
1546
+ # An unknown spelling is not licence to guess a tolerance.
1547
+ return (0.0 if a == b else float("inf")), 0.0
1548
+ rtol = _screen_rtol(canon)
1549
+ if rtol is None:
1550
+ return (0.0 if a == b else float("inf")), 0.0
1551
+ x = _as_float64(a, canon)
1552
+ y = _as_float64(b, canon)
1553
+ finite = np.isfinite(x) & np.isfinite(y)
1554
+ if not np.array_equal(np.isfinite(x), np.isfinite(y)):
1555
+ return float("inf"), float(rtol) # one produced NaN/Inf, one did not
1556
+ scale = float(np.abs(y[finite]).max()) if finite.any() else 0.0
1557
+ if scale == 0.0:
1558
+ return (0.0 if np.array_equal(x, y) else float("inf")), float(rtol)
1559
+ return float(np.abs(x[finite] - y[finite]).max() / scale), float(rtol)
1560
+
1561
+
1562
+ def configs_agreeing_with_oracle(results, oracle, dtype_name):
1563
+ """The configurations an oracle does not contradict, or None if there is none.
1564
+
1565
+ `screen_configs` decides by CONSENSUS among the candidates, which was the
1566
+ right answer to the 2026-09-07 incident: anchoring on a nominated reference
1567
+ inverts the moment that reference is the broken one. But consensus is a
1568
+ VOTE, and a vote has two failure modes it cannot see —
1569
+
1570
+ * the majority cluster is wrong in the same way, so the minority that is
1571
+ right gets excluded;
1572
+ * every candidate agrees and all are wrong, in which case the screen
1573
+ returns the whole space and says nothing.
1574
+
1575
+ Neither is hypothetical on hardware nobody has looked at. The certified
1576
+ directory covers `nvidia/volta` and nothing else, so on any other card the
1577
+ runtime sweeps and this screen is the only thing between the user and a
1578
+ wrong kernel — over a space that DIFFERS by target: the same flash tile
1579
+ needs 98 304 bytes on sm_70 and 164 352 on sm_86.
1580
+
1581
+ So when an oracle exists, it overrules the vote. When it does not, this
1582
+ returns None rather than pretending: the caller keeps the consensus it had
1583
+ and knows that is what it has. Refusing on an oracle nobody computed would
1584
+ be the same silence in the other direction.
1585
+
1586
+ `results` are `(config, output_bytes)` pairs; the comparison and its
1587
+ tolerance are the profile's own, the very ones the screen already applies
1588
+ between candidates.
1589
+ """
1590
+ if oracle is None:
1591
+ return None
1592
+ kept = []
1593
+ for entry in results:
1594
+ config, produced = entry[0], entry[1]
1595
+ if produced == oracle:
1596
+ kept.append(entry)
1597
+ continue
1598
+ deviation, tolerance = _deviation(produced, oracle, dtype_name)
1599
+ if deviation <= tolerance:
1600
+ kept.append(entry)
1601
+ return kept
1602
+
1603
+
1604
+ #: A callable `(tuner, key, buffers) -> list[bytes] | None` returning the
1605
+ #: CORRECT contents of the screened buffers, or None where none can be had.
1606
+ #: Default None: the screen keeps the consensus it has always kept, and this
1607
+ #: whole path costs one `is None`.
1608
+ #:
1609
+ #: `configs_agreeing_with_oracle` is the primitive and it takes ONE buffer; a
1610
+ #: screened result is a SNAPSHOT, a list of buffers with a dtype each. The
1611
+ #: primitive was correct and tested from the day it was written and still could
1612
+ #: not be called from here, because nothing carried a snapshot to it. That
1613
+ #: adapter is `_oracle_keeps` below, and its absence is why the overrule sat
1614
+ #: unwired: a helper whose every test passes can still have no seam.
1615
+ _SCREEN_ORACLE = None
1616
+
1617
+
1618
+ def set_screen_oracle(provider) -> None:
1619
+ """Install (or clear, with None) the oracle provider for the screen."""
1620
+ global _SCREEN_ORACLE
1621
+ _SCREEN_ORACLE = provider
1622
+
1623
+
1624
+ def _oracle_keeps(results, oracle, buffers):
1625
+ """The candidates that agree with the oracle on EVERY screened buffer."""
1626
+ kept = []
1627
+ for entry in results:
1628
+ agrees = True
1629
+ for (_a, _n, dtype_name), produced, reference in zip(buffers, entry[1], oracle):
1630
+ if not configs_agreeing_with_oracle([(entry[0], produced)],
1631
+ reference, dtype_name):
1632
+ agrees = False
1633
+ break
1634
+ if agrees:
1635
+ kept.append(entry)
1636
+ return kept
1637
+
1638
+
1639
+ def _make_agree(buffers):
1640
+ """The screen's agreement predicate for one set of screened buffers."""
1641
+ def agree(one, other):
1642
+ worst, tol, name = 0.0, 0.0, "?"
1643
+ for (_a, _n, dtype_name), x, y in zip(buffers, one, other):
1644
+ if x == y:
1645
+ continue
1646
+ deviation, tolerance = _deviation(x, y, dtype_name)
1647
+ if deviation > worst:
1648
+ worst, tol, name = deviation, tolerance, dtype_name
1649
+ return worst <= tol, worst, tol, name
1650
+ return agree
1651
+
1652
+
1653
+ def _cluster(results, agree):
1654
+ """Candidates grouped by mutual agreement, first-fit against each group."""
1655
+ clusters: List[list] = []
1656
+ for entry in results:
1657
+ for cluster in clusters:
1658
+ ok, _w, _t, _d = agree(entry[1], cluster[0][1])
1659
+ if ok:
1660
+ cluster.append(entry)
1661
+ break
1662
+ else:
1663
+ clusters.append([entry])
1664
+ return clusters
1665
+
1666
+
1667
+ def _largest_agreement(results, buffers):
1668
+ """What the VOTE would have seated — asked only to say whether it was
1669
+ about to be wrong. Never used to decide anything once an oracle exists."""
1670
+ clusters = _cluster(results, _make_agree(buffers))
1671
+ if not clusters:
1672
+ return []
1673
+ return max(clusters, key=len)
1674
+
1675
+
1676
+ #: what `do_bench` allocates on top of the arguments, once per call —
1677
+ #: an L2-sized scratch flushed before every timed launch.
1678
+ _BENCH_FLUSH_BYTES = 256 * 1024 * 1024
1679
+
1680
+
1681
+ def bench_would_swap(total_bytes: int):
1682
+ """(True, available_mb) when timing candidates would measure the swap.
1683
+
1684
+ The comparison carries NO TUNING MARGIN, and that is a different claim from
1685
+ "no constant at all" -- it used to say the latter, and stopped being true
1686
+ when `_BENCH_FLUSH_BYTES` was added (2026-09-17, after hat-s-x4 reached
1687
+ SIGKILL with the door open). What is added is not a cushion: it is the
1688
+ scratch `do_bench` REALLY ALLOCATES on every call, at the same name it
1689
+ allocates it with, so the two cannot drift. Nothing here is free to tune.
1690
+
1691
+ `available_mb` comes from `nbxtensor.backend.host_memory` -- our own authority for this
1692
+ quantity, never `Pages free` -- and is read live, never cached: pressure is
1693
+ a state of the moment, not of the process.
1694
+
1695
+ An unreadable platform (available_mb None) gates nothing and says nothing
1696
+ here: `memory_state` already names why it could not read.
1697
+ """
1698
+ try:
1699
+ from nbxtensor.backend import device_facts as _facts, host_memory as _hm # the library's own readings
1700
+ host_shares_memory_with_device, memory_state = _facts.host_shares_memory_with_device, _hm.memory_state
1701
+ if not host_shares_memory_with_device():
1702
+ return False, None # discrete memory: the host is not the bench
1703
+ avail_mb = memory_state().available_mb
1704
+ except Exception: # noqa: BLE001
1705
+ return False, None
1706
+ if avail_mb is None:
1707
+ return False, None
1708
+ # The sweep does not only hold the arguments: `do_bench` allocates a
1709
+ # 256 MiB L2 flush buffer per call, and Triton times one candidate after
1710
+ # another. Comparing the arguments alone therefore UNDERSTATES what the
1711
+ # sweep needs, which is how hat-s-x4 reached SIGKILL with the door open
1712
+ # (2026-09-17: a baddbmm key carrying 5.92 GB of arguments, the door
1713
+ # answering False, and the machine dying).
1714
+ need = total_bytes + _BENCH_FLUSH_BYTES
1715
+ import os as _os_bws
1716
+ if _os_bws.environ.get("NBX_BENCH_DOOR_DIAG"):
1717
+ _acct = ""
1718
+ try:
1719
+ from nbxtensor.memory.allocator import DeviceAllocator as _DA
1720
+ _live = _DA.memory_allocated() / 1e6
1721
+ _cached = sum(_DA._POOL.cached.values()) / 1e6
1722
+ _n = len(_DA._cuda_ptr_size)
1723
+ _acct = (f" | allocator: live {_live:.0f} MB, pooled {_cached:.0f} MB, "
1724
+ f"{_n} live pointers")
1725
+ except Exception as _e:
1726
+ _acct = f" | allocator accounting unavailable: {type(_e).__name__}"
1727
+ print(f"[BENCH_DOOR] args {total_bytes} + flush {_BENCH_FLUSH_BYTES} "
1728
+ f"= {need} bytes vs {avail_mb} MB available -> "
1729
+ f"{'REFUSE' if need > avail_mb * 2 ** 20 else 'allow'}{_acct}", flush=True)
1730
+ return need > avail_mb * 2 ** 20, avail_mb
1731
+
1732
+
1733
+ def autotune_shape_key(tuner, kwargs):
1734
+ """The key Triton's Autotuner will store this call's choice under.
1735
+
1736
+ Built exactly as `Autotuner.run` builds it — the `keys` arguments' values
1737
+ followed by every argument's dtype — from `tuner.nargs`, which holds the
1738
+ real arguments by the time `prune_configs` runs. Triton computes this key
1739
+ before calling `prune_configs` and does not pass it, so the screen rebuilt
1740
+ a DIFFERENT key (the constexpr kwargs) and recorded its unscreened seats
1741
+ under that one. The replay cache is written under Triton's key, so the
1742
+ `screened: false` mention never landed on any entry — seen on 2026-09-13
1743
+ by the production demonstration: 3 announcements, 0 stamped records. A
1744
+ key in one space, a record in another — the same family as the weight
1745
+ dict's two key spaces.
1746
+ """
1747
+ nargs = getattr(tuner, "nargs", None) or {}
1748
+ all_args = {**nargs, **(kwargs or {})}
1749
+ arg_names = getattr(tuner, "arg_names", None) or list(all_args)
1750
+ _args = {k: v for (k, v) in all_args.items() if k in arg_names}
1751
+ key = [_args[k] for k in (getattr(tuner, "keys", None) or []) if k in _args]
1752
+ for _, arg in _args.items():
1753
+ if hasattr(arg, "dtype"):
1754
+ key.append(str(arg.dtype))
1755
+ return tuple(key)
1756
+
1757
+
1758
+ def _call_screen_oracle(provider, tuner, key, buffers, meta):
1759
+ """Call the provider with the launch kwargs when its signature takes them."""
1760
+ import inspect
1761
+ try:
1762
+ n = len(inspect.signature(provider).parameters)
1763
+ except (TypeError, ValueError):
1764
+ n = 3
1765
+ if n >= 4:
1766
+ return provider(tuner, key, buffers, meta)
1767
+ return provider(tuner, key, buffers)
1768
+
1769
+
1770
+ def screen_configs(tuner, configs, key, meta=None, record_key=None):
1771
+ """Screen a key's candidates, then RELEASE what the sweep cached.
1772
+
1773
+ A sweep allocates: the candidates' own buffers, and a 256 MiB L2 flush
1774
+ buffer per `do_bench` call. The allocator pool retains freed blocks by
1775
+ design, and its release path is flush-on-OOM — which never runs on unified
1776
+ memory, because the OS KILLS the process instead of returning an allocation
1777
+ failure. So one key's sweep leaves its blocks resident and the next key
1778
+ starts poorer.
1779
+
1780
+ Measured on hat-s-x4 (M4 Pro, triton-ext, 2026-09-17), sampling available
1781
+ memory once a second while the run proceeded:
1782
+
1783
+ 18133 MB -> 15557 -> 15882 -> 9340 -> 3478 -> ... -> 210 MB, then SIGKILL
1784
+
1785
+ and the per-key door was RIGHT at each decision, because it looks at one key
1786
+ at a time:
1787
+
1788
+ args 5.92 GB + flush 0.27 GB vs 12695 MB available -> allow
1789
+ args 3.27 GB + flush 0.27 GB vs 8080 MB available -> allow
1790
+
1791
+ The door is not wrong; it is blind to accumulation. Releasing the pool when
1792
+ a key's sweep finishes is what makes each door's view of "available" true.
1793
+ """
1794
+ try:
1795
+ return _screen_configs(tuner, configs, key, meta=meta,
1796
+ record_key=record_key)
1797
+ finally:
1798
+ try:
1799
+ from nbxtensor.memory.allocator import DeviceAllocator
1800
+ DeviceAllocator.empty_cache_pool()
1801
+ except Exception: # noqa: BLE001
1802
+ pass # a sweep must not fail on cleanup
1803
+
1804
+
1805
+ def _screen_configs(tuner, configs, key, meta=None, record_key=None):
1806
+ """Run every candidate once and keep the ones that agree with each other.
1807
+
1808
+ `record_key` is the key the choice will be stored under (Triton's cache
1809
+ key, see `autotune_shape_key`); `key` is the screen's own de-duplication
1810
+ key. An unscreened seat is recorded under `record_key` so the replay cache
1811
+ can carry the mention.
1812
+
1813
+ Agreement is decided by CONSENSUS, not against a nominated reference.
1814
+ Anchoring on one config inverts the moment that config is the broken one:
1815
+ measured 2026-09-07, `matmul_kernel`'s first config was one of the three
1816
+ that wrote half the output, so screening against it excluded the seven
1817
+ correct ones. There is no way to know in advance which config is right —
1818
+ that is the whole problem — so the screen asks which answer the configs
1819
+ agree on, and treats the rest as the outliers they are.
1820
+
1821
+ A kernel writes into buffers the caller owns, so the screen snapshots
1822
+ them first and restores that snapshot before each run; otherwise an
1823
+ accumulating kernel would be compared against its own previous output.
1824
+
1825
+ Returns the configs in the consensus. Raises when there is no consensus —
1826
+ when the candidates split evenly — because then nothing here can tell
1827
+ which half is right, and choosing on speed would be choosing at random.
1828
+ """
1829
+ if len(configs) < 2:
1830
+ return configs
1831
+
1832
+ # De-duplicated by the SHAPE key, not the constexpr key: ten distinct
1833
+ # conv shapes share one constexpr tuple, and keyed by it the screen ran
1834
+ # on the first and silently skipped the other nine (2026-09-13, live).
1835
+ _rk = record_key if record_key is not None else key
1836
+ seen = _SCREEN_CACHE.setdefault(id(tuner), set())
1837
+ if _rk in seen:
1838
+ return configs
1839
+ seen.add(_rk)
1840
+
1841
+ kernel_name = getattr(tuner.base_fn, "__name__", str(tuner))
1842
+ named = dict(tuner.nargs or {})
1843
+ if not named:
1844
+ return _seat_unscreened(kernel_name, _rk, configs, len(configs),
1845
+ "the tuner carries no named arguments, so the "
1846
+ "screen has nothing to compare")
1847
+ args = [named[name] for name in tuner.arg_names if name in named]
1848
+ buffers = _writable_buffers(args)
1849
+ if buffers is None:
1850
+ return _seat_unscreened(kernel_name, _rk, configs, len(configs),
1851
+ "a strided view among the arguments, which the "
1852
+ "screen cannot snapshot")
1853
+ if not buffers:
1854
+ return _seat_unscreened(kernel_name, _rk, configs, len(configs),
1855
+ "no writable buffer to compare")
1856
+
1857
+ from nbxtensor.ops._configs import active_vendor_profile
1858
+
1859
+ budget = active_vendor_profile().get("autotune_screen_max_bytes")
1860
+ if budget is None:
1861
+ raise RuntimeError(
1862
+ "the hardware profile declares no `autotune_screen_max_bytes`: "
1863
+ "the screen will not decide for itself how much memory traffic a "
1864
+ "tuning step may cost")
1865
+ total = sum(nbytes for _a, nbytes, _d in buffers)
1866
+ if _needs_windowing(buffers, int(budget)) is not None:
1867
+ # Two doors meet here and both stay. Beyond the SCREEN budget the compare
1868
+ # is skipped (the Dell's ruling of 2026-09-12: the seat is announced as
1869
+ # UNSCREENED, never written to the certified directory). And Triton
1870
+ # would still time every candidate on these arguments; when they alone
1871
+ # exceed the machine's available memory the sweep measures the swap
1872
+ # (the Mac, 2026-09-13: a baddbmm key carrying 5.9 GB of arguments
1873
+ # against 4.5 GB available), so the sweep is cut to the single
1874
+ # first-declared config, the cut is SAID, and the choice is marked
1875
+ # unmeasured so capture() never persists it — recorded, it would
1876
+ # outlive the pressure that forced it.
1877
+ _swaps, _avail_mb = bench_would_swap(total)
1878
+ if _swaps:
1879
+ print(f"[AUTOTUNE_BENCH] "
1880
+ f"{getattr(tuner.base_fn, '__name__', tuner)}: arguments "
1881
+ f"total {total} bytes against {_avail_mb} MB available -- "
1882
+ f"timing candidates would measure the swap, not the "
1883
+ f"kernels. Taking the first declared config WITHOUT "
1884
+ f"measurement; the choice will not be persisted.",
1885
+ flush=True)
1886
+ from nbxtensor.ops import autotune_cache as _atc
1887
+ _atc.mark_unmeasured(tuner, key)
1888
+ return _seat_unscreened(
1889
+ kernel_name, _rk, configs[:1], 1,
1890
+ f"arguments total {total} bytes, over the profile's screening "
1891
+ f"budget {int(budget)} and over the {_avail_mb} MB available: "
1892
+ f"the sweep was cut to the first declared config, unmeasured")
1893
+ # OVER THE SCREENING BUDGET — screen on NAMED ROW WINDOWS rather than not at all.
1894
+ #
1895
+ # A fixed budget made the LARGEST shapes the LEAST verified ones, which is backwards:
1896
+ # the failure class that motivates screening a large shape is index overflow, and it
1897
+ # lives exactly where the budget stopped looking (the 2^31 addmm, 2026-09-23). The
1898
+ # budget is a rule about COST and the question is CORRECTNESS, so the cost is bounded
1899
+ # instead — a few row windows of the output, compared to the rack's row-windowed fp64
1900
+ # oracle, at a price set by the window and not by the shape (owner, 2026-09-23).
1901
+ #
1902
+ # The last window is anchored at the final row on purpose: an index that wraps shows
1903
+ # at the largest linear index or nowhere.
1904
+ windowed = _screen_on_windows(tuner, kernel_name, _rk, configs, args, meta,
1905
+ buffers, int(budget), total)
1906
+ if windowed is not None:
1907
+ return windowed
1908
+ return _seat_unscreened(
1909
+ kernel_name, _rk, configs, len(configs),
1910
+ f"arguments total {total} bytes, over the profile's screening "
1911
+ f"budget {int(budget)}, and the output could not be row-windowed")
1912
+
1913
+ before = _snapshot(buffers)
1914
+ meta = dict(meta or {})
1915
+
1916
+ # -- run each candidate once, from the same starting state --------------
1917
+ results, unrun = [], []
1918
+ for config in configs:
1919
+ try:
1920
+ _restore(buffers, before)
1921
+ tuner.fn.run(*args, **{**meta, **config.all_kwargs()})
1922
+ results.append((config, _snapshot(buffers)))
1923
+ except Exception:
1924
+ # A config the backend refuses is not a screen failure — the
1925
+ # autotuner already handles one that will not compile. Counted,
1926
+ # never swallowed: a screen that silently keeps what it could not
1927
+ # run is a screen that checked nothing.
1928
+ unrun.append(config)
1929
+ _restore(buffers, before)
1930
+
1931
+ if unrun:
1932
+ print(f"[AUTOTUNE_SCREEN] {kernel_name}: {len(unrun)} of "
1933
+ f"{len(configs)} configs could not be run for screening at key "
1934
+ f"{key}; they go to the timer unchecked", flush=True)
1935
+ if len(results) < 2:
1936
+ return _seat_unscreened(kernel_name, _rk, configs, len(results),
1937
+ "fewer than two candidates ran, so there is "
1938
+ "nothing to compare them against")
1939
+
1940
+ # -- an oracle, where one exists, OVERRULES the vote --------------------
1941
+ #
1942
+ # This is the whole point of the 2026-09-10 finding. Consensus has two
1943
+ # failure modes it cannot see: a majority wrong in the same way, and a
1944
+ # unanimous space that is wrong. Both are silent, and both are widest
1945
+ # exactly where we have never looked — outside `nvidia/volta`, where the
1946
+ # configuration space itself differs by target.
1947
+ oracle = None
1948
+ if _SCREEN_ORACLE is None:
1949
+ # A mechanism that is complete and switched off is the most expensive
1950
+ # form of a vacuous guard: it costs the price of writing it and returns
1951
+ # nothing. The only thing worse is one that is silent about being off.
1952
+ from nbxtensor.ops.screen_oracle import announce_no_oracle
1953
+ announce_no_oracle(kernel_name, key,
1954
+ why="no oracle provider is installed at all")
1955
+ if _SCREEN_ORACLE is not None:
1956
+ try:
1957
+ # The kernel's constexpr arguments (kernel_height, stride_*, padding_*,
1958
+ # groups, fp16 …) travel as launch KWARGS, not in `tuner.nargs`; a
1959
+ # provider that reads `nargs` alone never sees them. The GEMM oracles
1960
+ # need none and worked; the convolution oracle needs all of them and
1961
+ # returned None on every live key (2026-09-13, real-esrgan-x4 on
1962
+ # card 0) while its own suite was green. So the launch kwargs go to
1963
+ # the provider too, when it accepts them.
1964
+ # A refusal recorded for a PREVIOUS key is never read as this key's: the record is
1965
+ # per call, cleared before the provider speaks (a stale one surfaced as the reason
1966
+ # of an unrelated key, 2026-09-26 — a pre-existing pollution).
1967
+ from nbxtensor.ops.screen_oracle import set_last_refusal as _clear_refusal
1968
+ _clear_refusal(None)
1969
+ oracle = _call_screen_oracle(_SCREEN_ORACLE, tuner, key, buffers, meta)
1970
+ except Exception as exc: # an oracle that fails is not a launch failure
1971
+ print(f"[AUTOTUNE_ORACLE] {kernel_name}: the oracle provider raised "
1972
+ f"({type(exc).__name__}: {exc}); falling back to the consensus "
1973
+ f"at key {key}", flush=True)
1974
+ if oracle is not None:
1975
+ kept = _oracle_keeps(results, oracle, buffers)
1976
+ refused = [c for c, _ in results if c not in [k for k, _ in kept]]
1977
+ if not kept:
1978
+ raise RuntimeError(
1979
+ f"NeuroBrix autotune screen: {kernel_name} at key {key} — the "
1980
+ f"fp64 oracle contradicts EVERY candidate ({len(results)} of "
1981
+ f"{len(results)}). A consensus would have returned the whole "
1982
+ f"space and said nothing. Refusing to seat any of them.")
1983
+ if len(kept) < len(results):
1984
+ print(f"[AUTOTUNE_ORACLE] {kernel_name}: the oracle refused "
1985
+ f"{len(refused)} of {len(results)} configs at key {key}; "
1986
+ f"{len(kept)} kept. The vote is NOT consulted for this key.",
1987
+ flush=True)
1988
+ # Was the vote about to seat one of them? That is the finding.
1989
+ majority = _largest_agreement(results, buffers)
1990
+ wrong_majority = [c for c, _ in majority if c in refused]
1991
+ if wrong_majority:
1992
+ print(f"[AUTOTUNE_ORACLE] FINDING — {kernel_name} at key {key}: "
1993
+ f"{len(wrong_majority)} config(s) the CONSENSUS would have "
1994
+ f"seated are contradicted by the oracle. The majority was "
1995
+ f"wrong in the same way. This is the failure mode the "
1996
+ f"screen cannot see on its own.", flush=True)
1997
+ # The oracle adjudicated this key: the record says so, with the name of
1998
+ # what adjudicated it — the converse of the unscreened mention.
1999
+ _ADJUDICATED[(kernel_name, repr(_rk))] = "fp64 oracle"
2000
+ return [c for c, _ in kept] + unrun
2001
+
2002
+ # -- cluster by agreement ----------------------------------------------
2003
+ agree = _make_agree(buffers)
2004
+ clusters = _cluster(results, agree)
2005
+
2006
+ if len(clusters) == 1:
2007
+ return _seat_unscreened(
2008
+ kernel_name, _rk, [c for c, _ in results] + unrun, len(results),
2009
+ _no_oracle_reason(oracle)
2010
+ + ", and the candidates were unanimous — which is one of the two "
2011
+ "failure modes consensus cannot see")
2012
+
2013
+ clusters.sort(key=len, reverse=True)
2014
+ if len(clusters[0]) == len(clusters[1]):
2015
+ raise RuntimeError(
2016
+ f"NeuroBrix autotune screen: {kernel_name} at key {key} splits "
2017
+ f"into {len(clusters)} groups of configs that disagree with each "
2018
+ f"other, with no majority ("
2019
+ + ", ".join(str(len(c)) for c in clusters)
2020
+ + "). Refusing to choose one on speed.")
2021
+
2022
+ winners = {id(c) for c, _ in clusters[0]}
2023
+ dropped = []
2024
+ for cluster in clusters[1:]:
2025
+ for config, result in cluster:
2026
+ _ok, worst, tol, name = agree(result, clusters[0][0][1])
2027
+ dropped.append(ScreenedOut(kernel_name, key, str(config), name,
2028
+ worst, tol))
2029
+
2030
+ _SCREENED.extend(dropped)
2031
+ # A runtime exclusion that contradicts a CERTIFIED setting is a finding,
2032
+ # reported here and persisted — never a silence.
2033
+ try:
2034
+ from nbxtensor.ops.autotune_certified import report_contradictions
2035
+ report_contradictions(tuner, dropped)
2036
+ except Exception as exc: # the report must not turn into a launch failure
2037
+ print(f"[AUTOTUNE_SCREEN] could not check the certified directory: {exc}", flush=True)
2038
+ for entry in dropped:
2039
+ print(f"[AUTOTUNE_SCREEN] {entry.kernel}: config excluded before "
2040
+ f"timing — it disagrees with the {len(clusters[0])}-config "
2041
+ f"consensus by {entry.deviation:.3e} ({entry.dtype}, tolerance "
2042
+ f"{entry.tolerance:.1e}) at key {entry.key}: {entry.config}",
2043
+ flush=True)
2044
+ _record_screen_exclusions(dropped)
2045
+
2046
+ return _seat_unscreened(
2047
+ kernel_name, _rk, [c for c, _ in clusters[0]] + unrun, len(results),
2048
+ _no_oracle_reason(oracle)
2049
+ + f", and the winner is a {len(clusters[0])}-config majority — which "
2050
+ f"is the other failure mode consensus cannot see, a majority wrong "
2051
+ f"in the same way")
2052
+
2053
+
2054
+ def _record_screen_exclusions(dropped) -> None:
2055
+ """Hand the exclusions to the persisted sweep, so its readers see them."""
2056
+ try:
2057
+ from nbxtensor.ops import autotune_cache
2058
+ except Exception: # pragma: no cover
2059
+ return
2060
+ record = getattr(autotune_cache, "record_screen_exclusions", None)
2061
+ if record is None:
2062
+ return
2063
+ try:
2064
+ record([e._asdict() for e in dropped])
2065
+ except Exception as exc: # pragma: no cover
2066
+ print(f"[AUTOTUNE_SCREEN] could not persist exclusions: {exc}",
2067
+ flush=True)
2068
+
2069
+ # ---------------------------------------------------------------------------
2070
+ # The seam: route every `kernel[grid](...)` of the process through `launch`
2071
+ # ---------------------------------------------------------------------------
2072
+
2073
+ _installed = False
2074
+
2075
+
2076
+ #: The package whose kernels this launcher owns. A kernel defined anywhere else —
2077
+ #: torch's in-tree Triton ops, a third-party library's — is launched by Triton's own
2078
+ #: path: this engine's allocator never handed out its memory and has no claim on it.
2079
+ #: the packages whose kernels this launcher runs: the engine's, and the tensor library's (nbxtensor, where the kernels move)
2080
+ _OWN_PACKAGES = ("neurobrix", "nbxtensor")
2081
+
2082
+
2083
+ def _is_ours(jit_fn) -> bool:
2084
+ """True when this JITFunction was defined inside the engine or its tensor library.
2085
+
2086
+ Read from the function's own module, which is the authority: a kernel's home is
2087
+ where it was written, not a name we maintain a list of. An unreadable module is
2088
+ treated as foreign — the conservative direction, since refusing a foreign launch
2089
+ breaks a working library while passing one through only forgoes a check that was
2090
+ never ours to make.
2091
+ """
2092
+ fn = getattr(jit_fn, "fn", None)
2093
+ module = getattr(fn, "__module__", None) or getattr(jit_fn, "__module__", None)
2094
+ if not isinstance(module, str):
2095
+ return False
2096
+ return any(module == p or module.startswith(p + ".") for p in _OWN_PACKAGES)
2097
+
2098
+
2099
+ def install(force: Optional[bool] = None) -> bool:
2100
+ """Route `JITFunction.__getitem__` through the NeuroBrix launcher.
2101
+ `NBX_LAUNCHER=triton` keeps upstream's (the differential arm)."""
2102
+ global _installed
2103
+ if _installed:
2104
+ return True
2105
+ if force is None and os.environ.get("NBX_LAUNCHER", "nbx").lower() == "triton":
2106
+ return False
2107
+ try:
2108
+ from triton.runtime.jit import JITFunction
2109
+ except ModuleNotFoundError:
2110
+ # A compiled-only install (no Triton wheel: a Mac without it, a CPU
2111
+ # box) has no kernel launcher to route; the package import that
2112
+ # installs the seam must not be the import that ends the engine.
2113
+ # `tests/unit/cli/test_compiled_mode_needs_no_triton.py` was red on
2114
+ # this line from c8ed017 to 2026-09-16.
2115
+ return False
2116
+
2117
+ # The seam is process-wide: `JITFunction` is Triton's, so patching it routes
2118
+ # EVERY Triton kernel in the process — including ones this engine does not own.
2119
+ # torch 2.14 ships in-tree Triton implementations under `torch/_native/ops/`
2120
+ # (bmm_outer_product, foreach_mm, norm, polar, scatter_add, sum, topk) and
2121
+ # dispatches eager aten calls to them when a shape condition matches. Those
2122
+ # kernels are torch's, launched on torch's memory, and this launcher refused
2123
+ # them at its ownership rule — `aten.bmm::0`, parameter 'A_ptr', on the warm
2124
+ # compiled path of GLM-4.1V, Janus-Pro-7B and Sana-1600M (2026-09-17).
2125
+ #
2126
+ # The ownership rule is right and stays: a NeuroBrix kernel may not read memory
2127
+ # this engine did not hand out. It says nothing about a foreign library's kernel
2128
+ # on that library's own memory, and it is not ours to police. So the seam asks
2129
+ # whose kernel it is and gives a foreign one back to Triton's own path,
2130
+ # unwrapped — the differential arm `NBX_LAUNCHER=triton` already proves that
2131
+ # path works.
2132
+ _upstream_getitem = JITFunction.__getitem__
2133
+ _upstream_run = JITFunction.run
2134
+
2135
+ def __getitem__(self, grid):
2136
+ if not _is_ours(self):
2137
+ return _upstream_getitem(self, grid)
2138
+ return lambda *args, **kwargs: launch(self, grid, *args, **kwargs)
2139
+
2140
+ def run(self, *args, grid=None, warmup=False, **kwargs):
2141
+ """`JITFunction.run` is the Autotuner's launch path (every autotuned
2142
+ kernel: mm, bmm, addmm, conv2d, ...) and `warmup`'s — both routed here."""
2143
+ if not _is_ours(self):
2144
+ return _upstream_run(self, *args, grid=grid, warmup=warmup, **kwargs)
2145
+ if warmup:
2146
+ prepare(self, args, kwargs)
2147
+ return None
2148
+ return launch(self, grid, *args, **kwargs)
2149
+
2150
+ JITFunction.__getitem__ = __getitem__
2151
+ JITFunction.run = run
2152
+
2153
+ # The correctness screen, at the seam the autotuner itself uses to narrow
2154
+ # its candidates: `prune_configs` is called once per key, immediately
2155
+ # before anything is timed, and by then `self.nargs` holds the real
2156
+ # arguments. Wrapping it means the screen sees exactly the configs that
2157
+ # were about to be benchmarked, on exactly the tensors they would run on.
2158
+ from triton.runtime.autotuner import Autotuner
2159
+
2160
+ if not getattr(Autotuner.prune_configs, "_nbx_screened", False):
2161
+ _upstream_prune = Autotuner.prune_configs
2162
+
2163
+ def prune_configs(self, kwargs, *rest):
2164
+ configs = _upstream_prune(self, kwargs, *rest)
2165
+ if os.environ.get("NBX_AUTOTUNE_SCREEN", "on").lower() == "off":
2166
+ return configs
2167
+ key = tuple(sorted(
2168
+ (k, str(v)) for k, v in (kwargs or {}).items()
2169
+ if isinstance(v, (int, float, bool, str))))
2170
+ return screen_configs(self, list(configs), key, kwargs,
2171
+ record_key=autotune_shape_key(self, kwargs))
2172
+
2173
+ prune_configs._nbx_screened = True
2174
+ prune_configs._nbx_upstream = _upstream_prune
2175
+ Autotuner.prune_configs = prune_configs
2176
+
2177
+ # The screen consults an fp64 oracle by default. A CERTIFIED key never
2178
+ # reaches the screen, so this costs nothing where a certification exists;
2179
+ # it is paid only on an uncertified key, once for the key and not once per
2180
+ # candidate. `NBX_SCREEN_ORACLE=off` keeps the bare consensus for the
2181
+ # differential arm.
2182
+ if os.environ.get("NBX_SCREEN_ORACLE", "on").lower() != "off":
2183
+ try:
2184
+ from nbxtensor.ops.screen_oracle import install as _install_oracle
2185
+ _install_oracle()
2186
+ except Exception as exc: # never block a launch
2187
+ print(f"[AUTOTUNE_ORACLE] the oracle provider could not be "
2188
+ f"installed ({type(exc).__name__}: {exc}); the screen runs "
2189
+ f"on the consensus alone", flush=True)
2190
+
2191
+ # A backend refusal for ONE candidate config costs that config, not the
2192
+ # run. Triton's `_bench` already scores `OutOfResources` and friends `inf`
2193
+ # and carries on; `MetalNonRecoverableError` descends from `RuntimeError`,
2194
+ # so nothing catches it and it ends the sweep. Measured 2026-09-12: five
2195
+ # of six blocked models died inside the sweep on one refused config while
2196
+ # `conv2d_forward_kernel` declares eighteen, eleven of them servable.
2197
+ try:
2198
+ from nbxtensor.ops.autotune_refusals import install as _install_refusals
2199
+ if not _install_refusals():
2200
+ print("[AUTOTUNE_REFUSED] the autotuner is not importable; a "
2201
+ "refused config will end the sweep instead of being "
2202
+ "excluded", flush=True)
2203
+ except Exception as exc: # never block a launch
2204
+ print(f"[AUTOTUNE_REFUSED] the refusal policy could not be installed "
2205
+ f"({type(exc).__name__}: {exc}); a refused config will end the "
2206
+ f"sweep", flush=True)
2207
+
2208
+ _installed = True
2209
+ return True
2210
+
2211
+
2212
+ # ---------------------------------------------------------------------------
2213
+ # The benchmarker the autotuner uses: CUDA events through the allocator's
2214
+ # runtime handle — Triton's own asks torch for its timing and its buffers.
2215
+ # ---------------------------------------------------------------------------
2216
+
2217
+ def do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode="mean", budget_ms=None, **_):
2218
+ """Same contract as `triton.testing.do_bench` (milliseconds; quantiles or
2219
+ a mean/min/max), with the L2 flush and the timing done through the
2220
+ engine's runtime (`DeviceAllocator`: events on the legacy stream)."""
2221
+ import numpy as np
2222
+ from nbxtensor.memory.allocator import DeviceAllocator
2223
+ from nbxtensor.tensor import NBXTensor
2224
+ from nbxtensor.dtype.nbx_dtype import NBXDtype
2225
+ fn()
2226
+ DeviceAllocator.stream_synchronize(0)
2227
+ # an L2-sized scratch flushed before every timed call, as upstream does.
2228
+ # The SIZE is `_BENCH_FLUSH_BYTES`, not a literal repeated here:
2229
+ # `bench_would_swap` counts this allocation when it decides whether a sweep
2230
+ # would measure the swap, and a gate counting a different number from the
2231
+ # one actually allocated is a gate about nothing.
2232
+ flush = NBXTensor.empty((_BENCH_FLUSH_BYTES // 4,), dtype=NBXDtype.float32, device="cuda")
2233
+ # ONE probing launch before the four that finish the estimate. The
2234
+ # estimate used to be five blind launches -- so a pathologically slow
2235
+ # candidate (2.5 min a launch was measured on a bf16 matmul sweep) cost
2236
+ # five launches, >12 minutes, before any number existed. After a single
2237
+ # launch its time IS known, and a candidate whose one launch exceeds the
2238
+ # budget cannot win the sweep anyway; benching it further buys nothing.
2239
+ #
2240
+ # For a candidate under budget the arithmetic is unchanged: one launch
2241
+ # plus four launches, and t_est is the mean of the five -- the same
2242
+ # estimator, split across two event pairs instead of one.
2243
+ if budget_ms is None:
2244
+ # Triton's `_bench -> self.do_bench(kernel_call, quantiles=...)` chain
2245
+ # is not ours to re-sign, so the sweep wrapper PUBLISHES the budget and
2246
+ # this end consults it -- explicit coupling, documented at both ends.
2247
+ try:
2248
+ from nbxtensor.ops.autotune_refusals import current_budget_ms
2249
+ budget_ms = current_budget_ms()
2250
+ except Exception: # noqa: BLE001
2251
+ budget_ms = None
2252
+ t_probe = _time_ms(fn, 1)
2253
+ if budget_ms is not None and t_probe > budget_ms:
2254
+ raise CandidateOverTimeBudget(t_probe, budget_ms)
2255
+ try:
2256
+ from nbxtensor.ops.autotune_refusals import note_candidate_time
2257
+ note_candidate_time(t_probe)
2258
+ except Exception: # noqa: BLE001
2259
+ pass
2260
+ t_est = (t_probe + 4.0 * _time_ms(fn, 4)) / 5.0
2261
+ n_warmup = max(1, int(warmup / max(t_est, 1e-3)))
2262
+ n_repeat = max(1, int(rep / max(t_est, 1e-3)))
2263
+ for _i in range(n_warmup):
2264
+ fn()
2265
+ times = []
2266
+ for _i in range(n_repeat):
2267
+ DeviceAllocator.memset_cuda(flush.data_ptr(), 0, flush._nbytes)
2268
+ times.append(_time_ms(fn, 1))
2269
+ times = np.asarray(times, dtype=np.float64)
2270
+ if quantiles is not None:
2271
+ return [float(q) for q in np.quantile(times, quantiles)]
2272
+ return float(getattr(np, return_mode)(times)) if return_mode in ("mean", "min", "max", "median") else float(times.mean())
2273
+
2274
+
2275
+ class CandidateOverTimeBudget(Exception):
2276
+ """One autotune candidate exceeded the sweep's measured time budget.
2277
+
2278
+ Carries the two numbers so the exclusion can be ANNOUNCED with them: a
2279
+ candidate scored out for time without its time is a silent narrowing.
2280
+ """
2281
+
2282
+ def __init__(self, took_ms: float, budget_ms: float):
2283
+ self.took_ms = float(took_ms)
2284
+ self.budget_ms = float(budget_ms)
2285
+ super().__init__(
2286
+ f"one launch took {took_ms:.1f} ms against a budget of "
2287
+ f"{budget_ms:.1f} ms derived from this sweep's own measurements")
2288
+
2289
+
2290
+ def _time_ms(fn, n: int) -> float:
2291
+ """Wall time of n calls on the device, by CUDA events on the legacy stream."""
2292
+ from nbxtensor.memory.allocator import DeviceAllocator
2293
+ start = DeviceAllocator.create_event(timing=True); end = DeviceAllocator.create_event(timing=True)
2294
+ DeviceAllocator.record_event(start, 0)
2295
+ for _i in range(n):
2296
+ fn()
2297
+ DeviceAllocator.record_event(end, 0)
2298
+ DeviceAllocator.event_synchronize(end)
2299
+ ms = DeviceAllocator.event_elapsed_ms(start, end)
2300
+ DeviceAllocator.destroy_event(start); DeviceAllocator.destroy_event(end)
2301
+ return ms / n
2302
+
2303
+
2304
+ # ---------------------------------------------------------------------------
2305
+ # The launcher contract's CUDA client (`nbxtensor.backend.launcher_contract`,
2306
+ # Metal workstream 2026-09-05): a driver COMPILES a jit function plus an
2307
+ # explicit signature / constexprs / specialization markers into a
2308
+ # CompiledKernel, and LAUNCHES it with a grid and a flat argument list —
2309
+ # integer addresses for pointers (each verified with the allocator), typed
2310
+ # scalars otherwise. Streams and events are the allocator's. `CudaDriver`
2311
+ # carries this surface; the same checker that gates the Metal driver gates it.
2312
+ # ---------------------------------------------------------------------------
2313
+
2314
+ class CudaCompiledKernel:
2315
+ """The result of compiling once, launched many times (contract object)."""
2316
+
2317
+ binary_kind = "cubin"
2318
+
2319
+ def __init__(self, name, binary, block_size, shared_memory, constexprs, specialization, binding, signature, function):
2320
+ self.name = name
2321
+ self.binary = binary
2322
+ self.block_size = block_size
2323
+ self.shared_memory = shared_memory
2324
+ self.constexprs = dict(constexprs)
2325
+ self.specialization = dict(specialization)
2326
+ self.binding = tuple(binding)
2327
+ self._signature = signature # name -> triton type, non-constexpr, in binding order
2328
+ self._function = function
2329
+
2330
+ def launch(self, grid, args, stream: int = 0) -> None:
2331
+ args = list(args)
2332
+ if len(args) != len(self.binding):
2333
+ raise TypeError(f"NeuroBrix launcher: {self.name} takes {len(self.binding)} arguments, {len(args)} given")
2334
+ params = []
2335
+ for slot, value in zip(self.binding, args):
2336
+ if slot.is_pointer:
2337
+ params.append(("ptr", int(value))) # ownership is the driver's refusal, below
2338
+ else:
2339
+ params.append(_pack_param(slot.dtype, value))
2340
+ params.append(("ptr", 0)) # global scratch (Triton >= 3.6 ABI)
2341
+ params.append(("ptr", 0)) # profile scratch
2342
+ g = tuple(int(x) for x in grid) + (1,) * (3 - len(grid))
2343
+ CudaDriver.instance().launch(self._function, g, (self.block_size, 1, 1), self.shared_memory, int(stream), params)
2344
+
2345
+
2346
+ def driver() -> Driver:
2347
+ """The driver this process launches through — the launcher's backend
2348
+ and the contract's client are one object (`CudaDriver` on CUDA; the
2349
+ Metal driver registers itself behind the same interface)."""
2350
+ return active_driver()