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.
- nbxtensor/THIRD_PARTY_NOTICES.md +287 -0
- nbxtensor/__init__.py +17 -0
- nbxtensor/backend/__init__.py +1 -0
- nbxtensor/backend/capabilities.py +207 -0
- nbxtensor/backend/cpu/__init__.py +1 -0
- nbxtensor/backend/cpu/blas.py +102 -0
- nbxtensor/backend/cpu/pool.py +204 -0
- nbxtensor/backend/cpu_backend.py +162 -0
- nbxtensor/backend/device_facts.py +105 -0
- nbxtensor/backend/device_files.py +129 -0
- nbxtensor/backend/device_names.py +18 -0
- nbxtensor/backend/devices/amd/cdna.yml +228 -0
- nbxtensor/backend/devices/amd/cdna2.yml +228 -0
- nbxtensor/backend/devices/amd/cdna3.yml +227 -0
- nbxtensor/backend/devices/apple/apple_m1.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m1_max.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m1_pro.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m1_ultra.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m2.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m2_max.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m2_pro.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m2_ultra.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m3.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m3_max.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m3_pro.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m3_ultra.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m4.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m4_max.yml +266 -0
- nbxtensor/backend/devices/apple/apple_m4_pro.yml +308 -0
- nbxtensor/backend/devices/apple/apple_m5.yml +269 -0
- nbxtensor/backend/devices/apple/apple_m5_max.yml +269 -0
- nbxtensor/backend/devices/apple/apple_m5_pro.yml +269 -0
- nbxtensor/backend/devices/apple/apple_m5_ultra.yml +269 -0
- nbxtensor/backend/devices/apple/apple_m6.yml +269 -0
- nbxtensor/backend/devices/apple/apple_silicon.yml +263 -0
- nbxtensor/backend/devices/nvidia/ampere.yml +214 -0
- nbxtensor/backend/devices/nvidia/hopper.yml +213 -0
- nbxtensor/backend/devices/nvidia/volta.yml +411 -0
- nbxtensor/backend/faults.py +151 -0
- nbxtensor/backend/host_memory.py +185 -0
- nbxtensor/backend/launcher.py +2350 -0
- nbxtensor/backend/launcher_contract.py +429 -0
- nbxtensor/backend/metal_backend.py +661 -0
- nbxtensor/backend/metal_device.py +1274 -0
- nbxtensor/backend/process_memory.py +73 -0
- nbxtensor/backend/replay_graph.py +322 -0
- nbxtensor/backend/runtime.py +409 -0
- nbxtensor/backend/triton_ext_driver.py +735 -0
- nbxtensor/dtype/__init__.py +1 -0
- nbxtensor/dtype/convert.py +107 -0
- nbxtensor/dtype/dtypes.yml +70 -0
- nbxtensor/dtype/itemsize.py +166 -0
- nbxtensor/dtype/nbx_dtype.py +72 -0
- nbxtensor/dtype/operands.py +158 -0
- nbxtensor/dtype/outputs.py +64 -0
- nbxtensor/dtype/policy.py +155 -0
- nbxtensor/dtype/policy.yml +51 -0
- nbxtensor/dtype/promotion.py +112 -0
- nbxtensor/host.py +55 -0
- nbxtensor/io/__init__.py +1 -0
- nbxtensor/io/read_rate.py +197 -0
- nbxtensor/io/safetensors.py +314 -0
- nbxtensor/memory/__init__.py +9 -0
- nbxtensor/memory/allocator.py +1988 -0
- nbxtensor/memory/arena.py +75 -0
- nbxtensor/memory/device_memory.py +150 -0
- nbxtensor/memory/device_transfer.py +154 -0
- nbxtensor/memory/layout.py +43 -0
- nbxtensor/memory/memory_pool.py +191 -0
- nbxtensor/memory/metal.py +96 -0
- nbxtensor/memory/oom_advice.py +91 -0
- nbxtensor/memory/pointer_tables.py +323 -0
- nbxtensor/memory/pool.py +507 -0
- nbxtensor/memory/shadow.py +43 -0
- nbxtensor/memory/slab.py +273 -0
- nbxtensor/memory/trace_replay.py +122 -0
- nbxtensor/memory/traces/sana-video-2b-720p.16g.2026-11-02.tsv +13610 -0
- nbxtensor/memory/vmm.py +118 -0
- nbxtensor/ops/__init__.py +1 -0
- nbxtensor/ops/_autotune_policy.py +90 -0
- nbxtensor/ops/_common.py +32 -0
- nbxtensor/ops/_configs.py +1041 -0
- nbxtensor/ops/_cost.py +231 -0
- nbxtensor/ops/abs.py +17 -0
- nbxtensor/ops/adaptive_avg_pool2d.py +161 -0
- nbxtensor/ops/add.py +96 -0
- nbxtensor/ops/addcdiv.py +29 -0
- nbxtensor/ops/addcmul.py +29 -0
- nbxtensor/ops/addmv_op.py +64 -0
- nbxtensor/ops/addr.py +68 -0
- nbxtensor/ops/all_reduce.py +83 -0
- nbxtensor/ops/amax.py +40 -0
- nbxtensor/ops/any_reduce.py +83 -0
- nbxtensor/ops/arange_op.py +27 -0
- nbxtensor/ops/argmax.py +107 -0
- nbxtensor/ops/argmin.py +93 -0
- nbxtensor/ops/atan2.py +71 -0
- nbxtensor/ops/attention_wrappers.py +1648 -0
- nbxtensor/ops/autotune_bucket.py +137 -0
- nbxtensor/ops/autotune_cache.py +430 -0
- nbxtensor/ops/autotune_certified.py +1262 -0
- nbxtensor/ops/autotune_refusals.py +380 -0
- nbxtensor/ops/autotune_settings.py +58 -0
- nbxtensor/ops/avg_pool2d.py +73 -0
- nbxtensor/ops/baddbmm_op.py +132 -0
- nbxtensor/ops/batch_norm.py +163 -0
- nbxtensor/ops/binary_strided.py +196 -0
- nbxtensor/ops/bitwise_and.py +26 -0
- nbxtensor/ops/bitwise_left_shift.py +47 -0
- nbxtensor/ops/bitwise_not.py +24 -0
- nbxtensor/ops/bitwise_or.py +26 -0
- nbxtensor/ops/bitwise_right_shift.py +47 -0
- nbxtensor/ops/bucketize.py +37 -0
- nbxtensor/ops/cat_op.py +87 -0
- nbxtensor/ops/celu.py +20 -0
- nbxtensor/ops/clamp.py +24 -0
- nbxtensor/ops/clamp_min.py +22 -0
- nbxtensor/ops/conv2d.py +265 -0
- nbxtensor/ops/conv2d_m8n8k4.py +165 -0
- nbxtensor/ops/conv_depthwise2d.py +154 -0
- nbxtensor/ops/conv_transpose2d.py +148 -0
- nbxtensor/ops/conv_wrappers.py +1482 -0
- nbxtensor/ops/copy.py +18 -0
- nbxtensor/ops/copy_op.py +39 -0
- nbxtensor/ops/cos.py +17 -0
- nbxtensor/ops/count_nonzero.py +64 -0
- nbxtensor/ops/cross_entropy.py +67 -0
- nbxtensor/ops/cummax.py +109 -0
- nbxtensor/ops/cummin.py +105 -0
- nbxtensor/ops/cumsum.py +244 -0
- nbxtensor/ops/decode_attn_vec.py +162 -0
- nbxtensor/ops/depthwise_conv2d.py +132 -0
- nbxtensor/ops/dequant_gemv.py +190 -0
- nbxtensor/ops/dft.py +67 -0
- nbxtensor/ops/dispatch.py +1228 -0
- nbxtensor/ops/div.py +75 -0
- nbxtensor/ops/dot_op.py +67 -0
- nbxtensor/ops/dropout.py +23 -0
- nbxtensor/ops/dtype_convert.py +42 -0
- nbxtensor/ops/elementwise_launch.py +81 -0
- nbxtensor/ops/elementwise_wrappers.py +2315 -0
- nbxtensor/ops/elu.py +20 -0
- nbxtensor/ops/embedding.py +69 -0
- nbxtensor/ops/eq.py +32 -0
- nbxtensor/ops/erf.py +18 -0
- nbxtensor/ops/exp.py +17 -0
- nbxtensor/ops/exp2.py +18 -0
- nbxtensor/ops/fft_op.py +125 -0
- nbxtensor/ops/fill_op.py +20 -0
- nbxtensor/ops/flash_attention.py +307 -0
- nbxtensor/ops/flash_attention_m8n8k4.py +384 -0
- nbxtensor/ops/flash_decode.py +78 -0
- nbxtensor/ops/flip_op.py +60 -0
- nbxtensor/ops/floor.py +90 -0
- nbxtensor/ops/full_op.py +24 -0
- nbxtensor/ops/fused_moe.py +417 -0
- nbxtensor/ops/fused_moe_m8n8k4.py +99 -0
- nbxtensor/ops/fused_upsample_conv.py +413 -0
- nbxtensor/ops/gather_op.py +52 -0
- nbxtensor/ops/ge.py +32 -0
- nbxtensor/ops/gelu.py +31 -0
- nbxtensor/ops/gemv_vec.py +80 -0
- nbxtensor/ops/glu.py +43 -0
- nbxtensor/ops/grid_sampler.py +145 -0
- nbxtensor/ops/group_norm_tiled.py +67 -0
- nbxtensor/ops/groupnorm.py +123 -0
- nbxtensor/ops/gt.py +32 -0
- nbxtensor/ops/hardsigmoid.py +18 -0
- nbxtensor/ops/hardswish.py +19 -0
- nbxtensor/ops/im2col.py +57 -0
- nbxtensor/ops/index_add.py +95 -0
- nbxtensor/ops/index_put_op.py +92 -0
- nbxtensor/ops/index_select.py +121 -0
- nbxtensor/ops/index_wrappers.py +770 -0
- nbxtensor/ops/isfinite.py +20 -0
- nbxtensor/ops/isinf.py +20 -0
- nbxtensor/ops/isnan.py +18 -0
- nbxtensor/ops/launch_keys.py +1054 -0
- nbxtensor/ops/launch_observer.py +54 -0
- nbxtensor/ops/layernorm.py +71 -0
- nbxtensor/ops/le.py +35 -0
- nbxtensor/ops/leaky_relu.py +19 -0
- nbxtensor/ops/lerp.py +76 -0
- nbxtensor/ops/linalg_wrappers.py +1370 -0
- nbxtensor/ops/linspace_op.py +47 -0
- nbxtensor/ops/log.py +17 -0
- nbxtensor/ops/log_sigmoid.py +23 -0
- nbxtensor/ops/logical_and.py +26 -0
- nbxtensor/ops/logical_not.py +24 -0
- nbxtensor/ops/logical_or.py +26 -0
- nbxtensor/ops/logical_xor.py +37 -0
- nbxtensor/ops/loss_wrappers.py +282 -0
- nbxtensor/ops/lt.py +32 -0
- nbxtensor/ops/masked_fill.py +21 -0
- nbxtensor/ops/matmul.py +377 -0
- nbxtensor/ops/matmul_m8n8k4.py +235 -0
- nbxtensor/ops/max_pool2d.py +74 -0
- nbxtensor/ops/max_pool3d.py +93 -0
- nbxtensor/ops/max_reduce.py +91 -0
- nbxtensor/ops/maximum.py +18 -0
- nbxtensor/ops/mean.py +40 -0
- nbxtensor/ops/min_reduce.py +91 -0
- nbxtensor/ops/minimum.py +18 -0
- nbxtensor/ops/mish.py +22 -0
- nbxtensor/ops/mma_m8n8k4.py +152 -0
- nbxtensor/ops/moe_align.py +195 -0
- nbxtensor/ops/moe_decode_vec.py +374 -0
- nbxtensor/ops/moe_device_residency.py +141 -0
- nbxtensor/ops/moe_wrappers.py +345 -0
- nbxtensor/ops/mse_loss.py +79 -0
- nbxtensor/ops/mul.py +51 -0
- nbxtensor/ops/nan_to_num.py +23 -0
- nbxtensor/ops/ne.py +32 -0
- nbxtensor/ops/neg.py +17 -0
- nbxtensor/ops/nllloss.py +124 -0
- nbxtensor/ops/nonzero.py +33 -0
- nbxtensor/ops/nonzero_op.py +62 -0
- nbxtensor/ops/norm_wrappers.py +579 -0
- nbxtensor/ops/ones_op.py +20 -0
- nbxtensor/ops/oracles/__init__.py +6 -0
- nbxtensor/ops/oracles/conv2d_fp64.py +450 -0
- nbxtensor/ops/pad_op.py +159 -0
- nbxtensor/ops/paged_decode_attn.py +215 -0
- nbxtensor/ops/pixel_shuffle.py +144 -0
- nbxtensor/ops/pixel_unshuffle.py +72 -0
- nbxtensor/ops/pow.py +137 -0
- nbxtensor/ops/prod.py +81 -0
- nbxtensor/ops/rand_op.py +52 -0
- nbxtensor/ops/random_wrappers.py +366 -0
- nbxtensor/ops/reciprocal.py +17 -0
- nbxtensor/ops/reduction_wrappers.py +845 -0
- nbxtensor/ops/relu.py +18 -0
- nbxtensor/ops/remainder.py +46 -0
- nbxtensor/ops/repeat_interleave.py +48 -0
- nbxtensor/ops/repeat_op.py +132 -0
- nbxtensor/ops/residual_chain.py +374 -0
- nbxtensor/ops/rmsnorm.py +68 -0
- nbxtensor/ops/rope.py +142 -0
- nbxtensor/ops/rsqrt.py +17 -0
- nbxtensor/ops/scaled_softmax.py +96 -0
- nbxtensor/ops/scatter_op.py +139 -0
- nbxtensor/ops/screen_oracle.py +479 -0
- nbxtensor/ops/selu.py +22 -0
- nbxtensor/ops/shape_wrappers.py +1032 -0
- nbxtensor/ops/sigmoid.py +19 -0
- nbxtensor/ops/silu.py +27 -0
- nbxtensor/ops/sin.py +17 -0
- nbxtensor/ops/softmax.py +169 -0
- nbxtensor/ops/softplus.py +27 -0
- nbxtensor/ops/sort_op.py +304 -0
- nbxtensor/ops/spectral_wrappers.py +481 -0
- nbxtensor/ops/sqrt.py +17 -0
- nbxtensor/ops/stack_op.py +84 -0
- nbxtensor/ops/std.py +113 -0
- nbxtensor/ops/strided_copy.py +200 -0
- nbxtensor/ops/sub.py +35 -0
- nbxtensor/ops/sum.py +40 -0
- nbxtensor/ops/tan.py +19 -0
- nbxtensor/ops/tanh.py +21 -0
- nbxtensor/ops/threshold.py +19 -0
- nbxtensor/ops/tiling.yml +91 -0
- nbxtensor/ops/tiling_policy.py +43 -0
- nbxtensor/ops/tiling_sizes.py +370 -0
- nbxtensor/ops/topk.py +267 -0
- nbxtensor/ops/trace_op.py +64 -0
- nbxtensor/ops/tril.py +64 -0
- nbxtensor/ops/triu.py +62 -0
- nbxtensor/ops/unfold_op.py +37 -0
- nbxtensor/ops/unique_op.py +97 -0
- nbxtensor/ops/upsample_bilinear2d.py +106 -0
- nbxtensor/ops/upsample_nearest2d.py +52 -0
- nbxtensor/ops/upsample_nearest3d.py +61 -0
- nbxtensor/ops/var.py +136 -0
- nbxtensor/ops/vector_norm.py +216 -0
- nbxtensor/ops/weight_norm.py +104 -0
- nbxtensor/ops/where.py +39 -0
- nbxtensor/ops/wrappers.py +825 -0
- nbxtensor/ops/zeros_op.py +20 -0
- nbxtensor/quant/__init__.py +1 -0
- nbxtensor/quant/quantized_tensor.py +187 -0
- nbxtensor/random/__init__.py +1 -0
- nbxtensor/random/mt_randn.py +154 -0
- nbxtensor/random/rng_pin.py +59 -0
- nbxtensor/random/rng_stream.py +118 -0
- nbxtensor/random/seeded_draw.py +223 -0
- nbxtensor/shape/__init__.py +157 -0
- nbxtensor/tensor.py +1867 -0
- nbxtensor-0.0.1.dist-info/METADATA +379 -0
- nbxtensor-0.0.1.dist-info/RECORD +291 -0
- nbxtensor-0.0.1.dist-info/WHEEL +4 -0
- 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()
|