kernel-lens 1.2.0__tar.gz → 1.3.0__tar.gz
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.
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/PKG-INFO +128 -29
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/README.md +111 -26
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/__init__.py +1 -1
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/backends/builder.py +17 -3
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/backends/ort_gen.py +55 -20
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/backends/trt_gen.py +13 -13
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/compiler/ast_analyzer.py +54 -7
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/compiler/core.py +79 -7
- kernel_lens-1.3.0/kernel_lens/compiler/manifest.py +171 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/compiler/tracer.py +47 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/runtime/engine.py +252 -16
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens.egg-info/PKG-INFO +128 -29
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens.egg-info/requires.txt +3 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/pyproject.toml +27 -3
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_conformance.py +58 -0
- kernel_lens-1.2.0/kernel_lens/compiler/manifest.py +0 -80
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/LICENSE +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/backends/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/compiler/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/compiler/interaction.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/compiler/onnx_exporter.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/config.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInfer.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferImpl.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferLegacyDims.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferPlugin.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferPluginBase.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferPluginUtils.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferRuntime.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferRuntimeBase.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferRuntimeCommon.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferRuntimePlugin.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/include/tensorrt/NvInferVersion.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/runtime/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/utils/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/utils/deployment.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens/utils/env_check.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens.egg-info/SOURCES.txt +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens.egg-info/dependency_links.txt +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/kernel_lens.egg-info/top_level.txt +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/setup.cfg +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_alignment_guard.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_architectural_fixes.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_edge_cases.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_fp8_gemm.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_fused_cross_entropy.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_fused_residual_rms.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_inplace.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_inplace_dtypes.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_int4_dequant_gemm.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_mamba_selective_scan.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_research_kernels.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_standalone.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_standalone_launch.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_standalone_ptx.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_tricky_triton_kernels.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_triton_vs_torch_compile.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_trt.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.3.0}/tests/test_unit_cpu.py +0 -0
|
@@ -1,8 +1,8 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: kernel-lens
|
|
3
|
-
Version: 1.
|
|
4
|
-
Summary:
|
|
5
|
-
Author:
|
|
3
|
+
Version: 1.3.0
|
|
4
|
+
Summary: Compile PyTorch models with custom @triton.jit kernels into ONNX Runtime and TensorRT C++ plugins
|
|
5
|
+
Author-email: Justin Duc <justin.duc@galiad-research.com>
|
|
6
6
|
License: MIT License
|
|
7
7
|
|
|
8
8
|
Copyright (c) 2026 Duc Justin
|
|
@@ -25,11 +25,23 @@ License: MIT License
|
|
|
25
25
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
26
26
|
SOFTWARE.
|
|
27
27
|
|
|
28
|
+
Project-URL: Homepage, https://just1truc.github.io/kernel_lens/
|
|
29
|
+
Project-URL: Documentation, https://just1truc.github.io/kernel_lens/
|
|
30
|
+
Project-URL: Changelog, https://github.com/Just1truc/kernel_lens/blob/main/CHANGELOG.md
|
|
28
31
|
Project-URL: Source Code, https://github.com/Just1truc/kernel_lens
|
|
29
32
|
Project-URL: Bug Tracker, https://github.com/Just1truc/kernel_lens/issues
|
|
33
|
+
Project-URL: Paper, https://doi.org/10.5281/zenodo.23194523
|
|
34
|
+
Keywords: triton,tensorrt,onnxruntime,compiler,gpu,cuda,kernel,inference,deep-learning
|
|
35
|
+
Classifier: Development Status :: 4 - Beta
|
|
36
|
+
Classifier: Intended Audience :: Science/Research
|
|
30
37
|
Classifier: Programming Language :: Python :: 3
|
|
38
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
39
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
40
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
41
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
31
42
|
Classifier: License :: OSI Approved :: MIT License
|
|
32
43
|
Classifier: Operating System :: POSIX :: Linux
|
|
44
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA
|
|
33
45
|
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
34
46
|
Requires-Python: >=3.10
|
|
35
47
|
Description-Content-Type: text/markdown
|
|
@@ -46,6 +58,8 @@ Provides-Extra: all
|
|
|
46
58
|
Requires-Dist: onnxruntime-gpu>=1.17.0; extra == "all"
|
|
47
59
|
Requires-Dist: tensorrt<11.0.0,>=8.6.0; extra == "all"
|
|
48
60
|
Requires-Dist: tensorrt-cu12-libs<11.0.0; extra == "all"
|
|
61
|
+
Provides-Extra: bench
|
|
62
|
+
Requires-Dist: transformers>=4.40.0; extra == "bench"
|
|
49
63
|
Dynamic: license-file
|
|
50
64
|
|
|
51
65
|
# KernelLens
|
|
@@ -100,6 +114,12 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
|
|
|
100
114
|
- **Dynamic grids in C++** — launch-grid expressions are transpiled to C++
|
|
101
115
|
integer math over runtime shapes, so one plugin serves dynamic batch and
|
|
102
116
|
sequence lengths.
|
|
117
|
+
- **Shape-generic engines** — `kl.compile(..., dynamic_shapes=True)` exports the
|
|
118
|
+
leading dimension as a symbolic axis, so a single engine serves any size in a
|
|
119
|
+
declared range via a TensorRT optimization profile. Verified on Blackwell
|
|
120
|
+
(`sm_120a`) and Ampere (`sm_80`): one compile at batch 8 serving batches
|
|
121
|
+
1-16, bit-exact on both engines. Off by default, because a shape-specialised
|
|
122
|
+
engine is what the published benchmarks measure.
|
|
103
123
|
- **Zero-copy VRAM binding** — PyTorch `data_ptr()` values are registered
|
|
104
124
|
directly with each engine; no host staging copies.
|
|
105
125
|
- **Blackwell PTX ISA clamping** — a Triton-generated `.version 9.3` header is
|
|
@@ -113,11 +133,16 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
|
|
|
113
133
|
dynamic scalars, strided and dynamic inputs, loops, conditionals, nested
|
|
114
134
|
helpers, in-place mutation, 3-D grids, autotune, and fp16/bf16.
|
|
115
135
|
|
|
116
|
-
> **Honest scope.** KernelLens is research-grade
|
|
117
|
-
>
|
|
118
|
-
>
|
|
119
|
-
>
|
|
120
|
-
>
|
|
136
|
+
> **Honest scope.** KernelLens is research-grade, and the results are reported
|
|
137
|
+
> per configuration rather than blended. The composition is bit-exact against
|
|
138
|
+
> native Triton (`0.00e+00`) across the research kernels, the fourteen
|
|
139
|
+
> conformance patterns and the quantized GEMMs, on both backends; two
|
|
140
|
+
> multi-operator fusions land just above zero and are reported as such. Keeping
|
|
141
|
+
> the kernel pays when the baseline leaves the graph unfused and breaks even when
|
|
142
|
+
> it already fuses. The TensorRT-LLM backend is design-only (see
|
|
143
|
+
> [`docs/assets/data/a100/trtllm_backend_design.md`](docs/assets/data/a100/trtllm_backend_design.md)),
|
|
144
|
+
> and the quantized GEMMs are validated on ONNX Runtime only. Full methodology
|
|
145
|
+
> and every raw JSON are under [`docs/assets/data/`](docs/assets/data/).
|
|
121
146
|
|
|
122
147
|
## Installation
|
|
123
148
|
|
|
@@ -129,7 +154,16 @@ pip install kernel-lens[all] # both
|
|
|
129
154
|
```
|
|
130
155
|
|
|
131
156
|
**Requirements:** Linux x86-64, Python ≥ 3.10, CUDA ≥ 12.0, PyTorch ≥ 2.0,
|
|
132
|
-
Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`.
|
|
157
|
+
Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
|
|
158
|
+
generator emits `IPluginV2DynamicExt`).
|
|
159
|
+
|
|
160
|
+
**Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
|
|
161
|
+
`.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
|
|
162
|
+
compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
|
|
163
|
+
specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
|
|
164
|
+
`TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
|
|
165
|
+
above 2 GB from disk rather than from a memory buffer (opt-in; see the
|
|
166
|
+
[changelog](CHANGELOG.md)).
|
|
133
167
|
|
|
134
168
|
## Quickstart
|
|
135
169
|
|
|
@@ -168,26 +202,73 @@ compiled = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm",
|
|
|
168
202
|
backends=["onnx", "tensorrt"])
|
|
169
203
|
out_ort = compiled.run((x,), backend="onnx")
|
|
170
204
|
out_trt = compiled.run((x,), backend="tensorrt")
|
|
205
|
+
|
|
206
|
+
# One engine for any batch: mark the leading dim dynamic. dynamic_max bounds
|
|
207
|
+
# the optimization profile (default 4x the traced size).
|
|
208
|
+
serving = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm_dyn",
|
|
209
|
+
backends=["onnx", "tensorrt"], dynamic_shapes=True)
|
|
210
|
+
for b in (1, 2, 4, 8, 16): # all through the engine built at b=4
|
|
211
|
+
serving.run((torch.randn(b, 512, 4096, device="cuda"),), backend="tensorrt")
|
|
171
212
|
```
|
|
172
213
|
|
|
173
214
|
## Results at a glance
|
|
174
215
|
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
the
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
216
|
+
Two architectures: an **RTX 5060 Ti** (Blackwell, `sm_120a`) for the folding,
|
|
217
|
+
single-operator, end-to-end and quantized results, and an **A100-SXM4-80GB**
|
|
218
|
+
(Ampere, `sm_80`) for the conformance, parity and real-model results. Results are
|
|
219
|
+
reported per configuration and never blended. Full methodology, complete tables
|
|
220
|
+
and every raw JSON are on the
|
|
221
|
+
[results page](https://just1truc.github.io/kernel_lens/#results) and under
|
|
222
|
+
[`docs/assets/data/`](docs/assets/data/).
|
|
223
|
+
|
|
224
|
+
**The code you don't write.** The deployment is **nine lines of Python**; the C++
|
|
225
|
+
it replaces is **1,137 lines across six files** for one fused RMSNorm kernel on
|
|
226
|
+
both backends, measured with `tests/measure_codegen.py`.
|
|
227
|
+
|
|
228
|
+
**One engine, many shapes.** With `dynamic_shapes=True` a single compiled engine
|
|
229
|
+
serves any leading-dimension size in range. Measured with
|
|
230
|
+
`tests/test_conformance.py`, which compiles once and then runs every shape through
|
|
231
|
+
that engine — `0.00e+00` on both backends, on Blackwell and on Ampere:
|
|
232
|
+
|
|
233
|
+
| Case | Compiled at | Served |
|
|
234
|
+
| :--- | :--- | :--- |
|
|
235
|
+
| `dynamic_batch_sweep` | batch 8 | batches 1, 2, 4, 8, 16 |
|
|
236
|
+
| `dynamic_grid_sweep` | 524,288 elements | 65,536 / 131,072 / 524,288 / 1,048,576 |
|
|
237
|
+
|
|
238
|
+
**Real pretrained models.** Every decoder `LayerNorm` of a real OPT checkpoint
|
|
239
|
+
(125M → 6.7B) replaced by the Triton RMSNorm kernel, graded against
|
|
240
|
+
`torch.compile` — the baseline that *captures the same Triton kernels* rather
|
|
241
|
+
than reimplementing them. The result splits by precision:
|
|
242
|
+
|
|
243
|
+
- **fp32** — `torch.compile` barely improves on eager (1.02–1.12×) because it
|
|
244
|
+
leaves the operators around the kernel unfused. KernelLens fuses them:
|
|
245
|
+
**3.18–5.45× faster** under TensorRT and **2.27–4.62×** under ONNX Runtime,
|
|
246
|
+
and the margin grows with model size.
|
|
247
|
+
- **fp16** — `torch.compile` optimises well there, so we only match it: TensorRT
|
|
248
|
+
edges it by **1.06–1.22×**, while ONNX Runtime is consistently **slower
|
|
249
|
+
(0.67–0.84×)** — per-node dispatch dominates when the graph already fuses.
|
|
250
|
+
- **Accuracy tracks precision, not speed.** fp32 holds the same greedy token on
|
|
251
|
+
**≥ 98.8%** of positions; fp16 under TensorRT falls to 93.4–97.9%. OPT-350M in
|
|
252
|
+
fp16 under ONNX Runtime is bit-exact.
|
|
253
|
+
|
|
254
|
+
**Fold + keep, and it pays.** On a two-layer block whose attention is a custom
|
|
255
|
+
squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
|
|
256
|
+
for the folded-only `torch.compile` formulation and **28.04 ms** for eager
|
|
257
|
+
folding — **1.99×** and **3.68×** faster. Against the unfused kernel path it is
|
|
258
|
+
**1.08×** faster than eager Triton (8.24 ms) and level with `torch.compile`
|
|
259
|
+
(7.55 ms). The engine folds the standard layers; the kernel keeps the operations
|
|
260
|
+
folding cannot express.
|
|
261
|
+
|
|
262
|
+
**Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
|
|
263
|
+
PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and the
|
|
264
|
+
quantized GEMMs, on both backends — and across all **fourteen** module patterns
|
|
265
|
+
of `tests/test_conformance.py`. Two multi-operator fusions are near-exact rather
|
|
266
|
+
than exact (`≤ 9.5e-07` residual+RMSNorm, `4.8e-07` Mamba's scan) and are
|
|
267
|
+
reported as such.
|
|
268
|
+
|
|
269
|
+
**Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
|
|
270
|
+
and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
|
|
271
|
+
7.85/7.92 ms).
|
|
191
272
|
|
|
192
273
|
## Documentation
|
|
193
274
|
|
|
@@ -195,9 +276,11 @@ the [results page](https://just1truc.github.io/kernel_lens/#results).
|
|
|
195
276
|
| :--- | :--- |
|
|
196
277
|
| [Problem](https://just1truc.github.io/kernel_lens/#problem) | The five gaps a naive export hits |
|
|
197
278
|
| [Pipeline](https://just1truc.github.io/kernel_lens/#pipeline) | What you build by hand vs. the compile/run pipeline |
|
|
198
|
-
| [Results](https://just1truc.github.io/kernel_lens/#results) | Folding + kernels, parity,
|
|
279
|
+
| [Results](https://just1truc.github.io/kernel_lens/#results) | Folding + kernels, parity, conformance, single-operator and end-to-end latency, **real pretrained models** |
|
|
280
|
+
| [Why an engine](https://just1truc.github.io/kernel_lens/#why) | `torch.compile` vs TensorRT, both directions |
|
|
199
281
|
| [Install & API](https://just1truc.github.io/kernel_lens/#install) | Installation and API reference |
|
|
200
282
|
| [Status](https://just1truc.github.io/kernel_lens/#status) | Roadmap |
|
|
283
|
+
| [FAQ](https://just1truc.github.io/kernel_lens/#faq) | Fair objections |
|
|
201
284
|
| [Paper](https://just1truc.github.io/kernel_lens/#paper) | Abstract and citation |
|
|
202
285
|
|
|
203
286
|
## Reproduce the results
|
|
@@ -212,14 +295,30 @@ python tests/test_research_kernels.py # numerical parity
|
|
|
212
295
|
python tests/benchmark_research_kernels.py # single-operator latency
|
|
213
296
|
python tests/benchmark_e2e_llama.py # end-to-end decoder
|
|
214
297
|
python tests/test_fp8_gemm.py # FP8 Tensor Core path
|
|
298
|
+
python tests/measure_codegen.py # generated C++ volume, per file
|
|
299
|
+
python tests/benchmark_real_model.py --model facebook/opt-1.3b # real OPT, both precisions
|
|
215
300
|
```
|
|
216
301
|
|
|
302
|
+
The real-model matrix (OPT-125M → 6.7B, fp32 and fp16) runs unattended on a fresh
|
|
303
|
+
CUDA host: `scripts/remote_setup.sh` provisions nvcc, TensorRT, `onnxruntime-gpu`
|
|
304
|
+
and `transformers` and installs the package; `scripts/remote_bench_all.sh` then
|
|
305
|
+
sweeps the matrix into `results/` and writes `summary.json`, recording failures
|
|
306
|
+
instead of aborting. See
|
|
307
|
+
[`docs/assets/data/a100/runpod_runbook.md`](docs/assets/data/a100/runpod_runbook.md) for the
|
|
308
|
+
exact invocation and the environment traps (CUDA 12.8's `nvcc` does not build
|
|
309
|
+
against a modern libstdc++; CUDA 13.3's does).
|
|
310
|
+
|
|
217
311
|
Raw measurement JSON is committed under [`docs/assets/data/`](docs/assets/data/).
|
|
218
312
|
|
|
219
313
|
## Technical report
|
|
220
314
|
|
|
221
|
-
The
|
|
222
|
-
|
|
315
|
+
The paper source is [`kernel_lens_fresh.tex`](kernel_lens_fresh.tex). Two derived
|
|
316
|
+
variants are generated from it and must not be edited by hand:
|
|
317
|
+
`scripts/make_mlsys.py` produces the 2-column MLSys form
|
|
318
|
+
(`kernel_lens_mlsys.tex`) and `scripts/make_tmlr_fresh.py` produces the TMLR
|
|
319
|
+
forms — `kernel_lens_tmlr_fresh.tex` (named preprint) and
|
|
320
|
+
`kernel_lens_tmlr_submission.tex` (anonymized; **do not publish it alongside the
|
|
321
|
+
named preprint**).
|
|
223
322
|
|
|
224
323
|
```bibtex
|
|
225
324
|
@article{duc2026kernellens,
|
|
@@ -227,7 +326,7 @@ variant is generated from it with `python3 scripts/make_tmlr.py`.
|
|
|
227
326
|
author = {Duc, Justin},
|
|
228
327
|
institution = {Galiad Research},
|
|
229
328
|
year = {2026},
|
|
230
|
-
version = {1.1
|
|
329
|
+
version = {1.2.1},
|
|
231
330
|
doi = {10.5281/zenodo.23194523},
|
|
232
331
|
url = {https://doi.org/10.5281/zenodo.23194523}
|
|
233
332
|
}
|
|
@@ -50,6 +50,12 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
|
|
|
50
50
|
- **Dynamic grids in C++** — launch-grid expressions are transpiled to C++
|
|
51
51
|
integer math over runtime shapes, so one plugin serves dynamic batch and
|
|
52
52
|
sequence lengths.
|
|
53
|
+
- **Shape-generic engines** — `kl.compile(..., dynamic_shapes=True)` exports the
|
|
54
|
+
leading dimension as a symbolic axis, so a single engine serves any size in a
|
|
55
|
+
declared range via a TensorRT optimization profile. Verified on Blackwell
|
|
56
|
+
(`sm_120a`) and Ampere (`sm_80`): one compile at batch 8 serving batches
|
|
57
|
+
1-16, bit-exact on both engines. Off by default, because a shape-specialised
|
|
58
|
+
engine is what the published benchmarks measure.
|
|
53
59
|
- **Zero-copy VRAM binding** — PyTorch `data_ptr()` values are registered
|
|
54
60
|
directly with each engine; no host staging copies.
|
|
55
61
|
- **Blackwell PTX ISA clamping** — a Triton-generated `.version 9.3` header is
|
|
@@ -63,11 +69,16 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
|
|
|
63
69
|
dynamic scalars, strided and dynamic inputs, loops, conditionals, nested
|
|
64
70
|
helpers, in-place mutation, 3-D grids, autotune, and fp16/bf16.
|
|
65
71
|
|
|
66
|
-
> **Honest scope.** KernelLens is research-grade
|
|
67
|
-
>
|
|
68
|
-
>
|
|
69
|
-
>
|
|
70
|
-
>
|
|
72
|
+
> **Honest scope.** KernelLens is research-grade, and the results are reported
|
|
73
|
+
> per configuration rather than blended. The composition is bit-exact against
|
|
74
|
+
> native Triton (`0.00e+00`) across the research kernels, the fourteen
|
|
75
|
+
> conformance patterns and the quantized GEMMs, on both backends; two
|
|
76
|
+
> multi-operator fusions land just above zero and are reported as such. Keeping
|
|
77
|
+
> the kernel pays when the baseline leaves the graph unfused and breaks even when
|
|
78
|
+
> it already fuses. The TensorRT-LLM backend is design-only (see
|
|
79
|
+
> [`docs/assets/data/a100/trtllm_backend_design.md`](docs/assets/data/a100/trtllm_backend_design.md)),
|
|
80
|
+
> and the quantized GEMMs are validated on ONNX Runtime only. Full methodology
|
|
81
|
+
> and every raw JSON are under [`docs/assets/data/`](docs/assets/data/).
|
|
71
82
|
|
|
72
83
|
## Installation
|
|
73
84
|
|
|
@@ -79,7 +90,16 @@ pip install kernel-lens[all] # both
|
|
|
79
90
|
```
|
|
80
91
|
|
|
81
92
|
**Requirements:** Linux x86-64, Python ≥ 3.10, CUDA ≥ 12.0, PyTorch ≥ 2.0,
|
|
82
|
-
Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`.
|
|
93
|
+
Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
|
|
94
|
+
generator emits `IPluginV2DynamicExt`).
|
|
95
|
+
|
|
96
|
+
**Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
|
|
97
|
+
`.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
|
|
98
|
+
compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
|
|
99
|
+
specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
|
|
100
|
+
`TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
|
|
101
|
+
above 2 GB from disk rather than from a memory buffer (opt-in; see the
|
|
102
|
+
[changelog](CHANGELOG.md)).
|
|
83
103
|
|
|
84
104
|
## Quickstart
|
|
85
105
|
|
|
@@ -118,26 +138,73 @@ compiled = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm",
|
|
|
118
138
|
backends=["onnx", "tensorrt"])
|
|
119
139
|
out_ort = compiled.run((x,), backend="onnx")
|
|
120
140
|
out_trt = compiled.run((x,), backend="tensorrt")
|
|
141
|
+
|
|
142
|
+
# One engine for any batch: mark the leading dim dynamic. dynamic_max bounds
|
|
143
|
+
# the optimization profile (default 4x the traced size).
|
|
144
|
+
serving = kl.compile(RMSNorm(4096).cuda(), (x,), name="llama3_rmsnorm_dyn",
|
|
145
|
+
backends=["onnx", "tensorrt"], dynamic_shapes=True)
|
|
146
|
+
for b in (1, 2, 4, 8, 16): # all through the engine built at b=4
|
|
147
|
+
serving.run((torch.randn(b, 512, 4096, device="cuda"),), backend="tensorrt")
|
|
121
148
|
```
|
|
122
149
|
|
|
123
150
|
## Results at a glance
|
|
124
151
|
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
the
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
152
|
+
Two architectures: an **RTX 5060 Ti** (Blackwell, `sm_120a`) for the folding,
|
|
153
|
+
single-operator, end-to-end and quantized results, and an **A100-SXM4-80GB**
|
|
154
|
+
(Ampere, `sm_80`) for the conformance, parity and real-model results. Results are
|
|
155
|
+
reported per configuration and never blended. Full methodology, complete tables
|
|
156
|
+
and every raw JSON are on the
|
|
157
|
+
[results page](https://just1truc.github.io/kernel_lens/#results) and under
|
|
158
|
+
[`docs/assets/data/`](docs/assets/data/).
|
|
159
|
+
|
|
160
|
+
**The code you don't write.** The deployment is **nine lines of Python**; the C++
|
|
161
|
+
it replaces is **1,137 lines across six files** for one fused RMSNorm kernel on
|
|
162
|
+
both backends, measured with `tests/measure_codegen.py`.
|
|
163
|
+
|
|
164
|
+
**One engine, many shapes.** With `dynamic_shapes=True` a single compiled engine
|
|
165
|
+
serves any leading-dimension size in range. Measured with
|
|
166
|
+
`tests/test_conformance.py`, which compiles once and then runs every shape through
|
|
167
|
+
that engine — `0.00e+00` on both backends, on Blackwell and on Ampere:
|
|
168
|
+
|
|
169
|
+
| Case | Compiled at | Served |
|
|
170
|
+
| :--- | :--- | :--- |
|
|
171
|
+
| `dynamic_batch_sweep` | batch 8 | batches 1, 2, 4, 8, 16 |
|
|
172
|
+
| `dynamic_grid_sweep` | 524,288 elements | 65,536 / 131,072 / 524,288 / 1,048,576 |
|
|
173
|
+
|
|
174
|
+
**Real pretrained models.** Every decoder `LayerNorm` of a real OPT checkpoint
|
|
175
|
+
(125M → 6.7B) replaced by the Triton RMSNorm kernel, graded against
|
|
176
|
+
`torch.compile` — the baseline that *captures the same Triton kernels* rather
|
|
177
|
+
than reimplementing them. The result splits by precision:
|
|
178
|
+
|
|
179
|
+
- **fp32** — `torch.compile` barely improves on eager (1.02–1.12×) because it
|
|
180
|
+
leaves the operators around the kernel unfused. KernelLens fuses them:
|
|
181
|
+
**3.18–5.45× faster** under TensorRT and **2.27–4.62×** under ONNX Runtime,
|
|
182
|
+
and the margin grows with model size.
|
|
183
|
+
- **fp16** — `torch.compile` optimises well there, so we only match it: TensorRT
|
|
184
|
+
edges it by **1.06–1.22×**, while ONNX Runtime is consistently **slower
|
|
185
|
+
(0.67–0.84×)** — per-node dispatch dominates when the graph already fuses.
|
|
186
|
+
- **Accuracy tracks precision, not speed.** fp32 holds the same greedy token on
|
|
187
|
+
**≥ 98.8%** of positions; fp16 under TensorRT falls to 93.4–97.9%. OPT-350M in
|
|
188
|
+
fp16 under ONNX Runtime is bit-exact.
|
|
189
|
+
|
|
190
|
+
**Fold + keep, and it pays.** On a two-layer block whose attention is a custom
|
|
191
|
+
squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
|
|
192
|
+
for the folded-only `torch.compile` formulation and **28.04 ms** for eager
|
|
193
|
+
folding — **1.99×** and **3.68×** faster. Against the unfused kernel path it is
|
|
194
|
+
**1.08×** faster than eager Triton (8.24 ms) and level with `torch.compile`
|
|
195
|
+
(7.55 ms). The engine folds the standard layers; the kernel keeps the operations
|
|
196
|
+
folding cannot express.
|
|
197
|
+
|
|
198
|
+
**Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
|
|
199
|
+
PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and the
|
|
200
|
+
quantized GEMMs, on both backends — and across all **fourteen** module patterns
|
|
201
|
+
of `tests/test_conformance.py`. Two multi-operator fusions are near-exact rather
|
|
202
|
+
than exact (`≤ 9.5e-07` residual+RMSNorm, `4.8e-07` Mamba's scan) and are
|
|
203
|
+
reported as such.
|
|
204
|
+
|
|
205
|
+
**Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
|
|
206
|
+
and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
|
|
207
|
+
7.85/7.92 ms).
|
|
141
208
|
|
|
142
209
|
## Documentation
|
|
143
210
|
|
|
@@ -145,9 +212,11 @@ the [results page](https://just1truc.github.io/kernel_lens/#results).
|
|
|
145
212
|
| :--- | :--- |
|
|
146
213
|
| [Problem](https://just1truc.github.io/kernel_lens/#problem) | The five gaps a naive export hits |
|
|
147
214
|
| [Pipeline](https://just1truc.github.io/kernel_lens/#pipeline) | What you build by hand vs. the compile/run pipeline |
|
|
148
|
-
| [Results](https://just1truc.github.io/kernel_lens/#results) | Folding + kernels, parity,
|
|
215
|
+
| [Results](https://just1truc.github.io/kernel_lens/#results) | Folding + kernels, parity, conformance, single-operator and end-to-end latency, **real pretrained models** |
|
|
216
|
+
| [Why an engine](https://just1truc.github.io/kernel_lens/#why) | `torch.compile` vs TensorRT, both directions |
|
|
149
217
|
| [Install & API](https://just1truc.github.io/kernel_lens/#install) | Installation and API reference |
|
|
150
218
|
| [Status](https://just1truc.github.io/kernel_lens/#status) | Roadmap |
|
|
219
|
+
| [FAQ](https://just1truc.github.io/kernel_lens/#faq) | Fair objections |
|
|
151
220
|
| [Paper](https://just1truc.github.io/kernel_lens/#paper) | Abstract and citation |
|
|
152
221
|
|
|
153
222
|
## Reproduce the results
|
|
@@ -162,14 +231,30 @@ python tests/test_research_kernels.py # numerical parity
|
|
|
162
231
|
python tests/benchmark_research_kernels.py # single-operator latency
|
|
163
232
|
python tests/benchmark_e2e_llama.py # end-to-end decoder
|
|
164
233
|
python tests/test_fp8_gemm.py # FP8 Tensor Core path
|
|
234
|
+
python tests/measure_codegen.py # generated C++ volume, per file
|
|
235
|
+
python tests/benchmark_real_model.py --model facebook/opt-1.3b # real OPT, both precisions
|
|
165
236
|
```
|
|
166
237
|
|
|
238
|
+
The real-model matrix (OPT-125M → 6.7B, fp32 and fp16) runs unattended on a fresh
|
|
239
|
+
CUDA host: `scripts/remote_setup.sh` provisions nvcc, TensorRT, `onnxruntime-gpu`
|
|
240
|
+
and `transformers` and installs the package; `scripts/remote_bench_all.sh` then
|
|
241
|
+
sweeps the matrix into `results/` and writes `summary.json`, recording failures
|
|
242
|
+
instead of aborting. See
|
|
243
|
+
[`docs/assets/data/a100/runpod_runbook.md`](docs/assets/data/a100/runpod_runbook.md) for the
|
|
244
|
+
exact invocation and the environment traps (CUDA 12.8's `nvcc` does not build
|
|
245
|
+
against a modern libstdc++; CUDA 13.3's does).
|
|
246
|
+
|
|
167
247
|
Raw measurement JSON is committed under [`docs/assets/data/`](docs/assets/data/).
|
|
168
248
|
|
|
169
249
|
## Technical report
|
|
170
250
|
|
|
171
|
-
The
|
|
172
|
-
|
|
251
|
+
The paper source is [`kernel_lens_fresh.tex`](kernel_lens_fresh.tex). Two derived
|
|
252
|
+
variants are generated from it and must not be edited by hand:
|
|
253
|
+
`scripts/make_mlsys.py` produces the 2-column MLSys form
|
|
254
|
+
(`kernel_lens_mlsys.tex`) and `scripts/make_tmlr_fresh.py` produces the TMLR
|
|
255
|
+
forms — `kernel_lens_tmlr_fresh.tex` (named preprint) and
|
|
256
|
+
`kernel_lens_tmlr_submission.tex` (anonymized; **do not publish it alongside the
|
|
257
|
+
named preprint**).
|
|
173
258
|
|
|
174
259
|
```bibtex
|
|
175
260
|
@article{duc2026kernellens,
|
|
@@ -177,7 +262,7 @@ variant is generated from it with `python3 scripts/make_tmlr.py`.
|
|
|
177
262
|
author = {Duc, Justin},
|
|
178
263
|
institution = {Galiad Research},
|
|
179
264
|
year = {2026},
|
|
180
|
-
version = {1.1
|
|
265
|
+
version = {1.2.1},
|
|
181
266
|
doi = {10.5281/zenodo.23194523},
|
|
182
267
|
url = {https://doi.org/10.5281/zenodo.23194523}
|
|
183
268
|
}
|
|
@@ -6,6 +6,20 @@ import shutil
|
|
|
6
6
|
|
|
7
7
|
from ..utils.env_check import find_nvcc
|
|
8
8
|
|
|
9
|
+
|
|
10
|
+
def _host_compiler():
|
|
11
|
+
"""Host compiler for nvcc.
|
|
12
|
+
|
|
13
|
+
CUDA toolkits before 12.9 declare rsqrt/sinpi/... without noexcept, which
|
|
14
|
+
GCC 13 rejects ("has a different exception specifier"), so the right choice
|
|
15
|
+
is environment-specific. Override with KERNEL_LENS_CCBIN.
|
|
16
|
+
"""
|
|
17
|
+
override = os.environ.get("KERNEL_LENS_CCBIN")
|
|
18
|
+
if override:
|
|
19
|
+
return override
|
|
20
|
+
return "g++-13" if os.path.exists("/usr/bin/g++-13") else "g++"
|
|
21
|
+
|
|
22
|
+
|
|
9
23
|
def get_cuda_home() -> str:
|
|
10
24
|
nvcc_path = find_nvcc()
|
|
11
25
|
if nvcc_path:
|
|
@@ -152,7 +166,7 @@ def build_ort_plugin(ort_plugins_dir: str, cache_dir: str):
|
|
|
152
166
|
arch_flag = "-gencode=arch=compute_75,code=sm_75"
|
|
153
167
|
|
|
154
168
|
user_trt_inc = os.path.expanduser("~/tensorrt_headers")
|
|
155
|
-
cpp_compiler =
|
|
169
|
+
cpp_compiler = _host_compiler()
|
|
156
170
|
ccbin_flag = ["-ccbin", cpp_compiler] if os.path.exists(f"/usr/bin/{cpp_compiler}") else []
|
|
157
171
|
|
|
158
172
|
cmd = [
|
|
@@ -318,7 +332,7 @@ def build_trt_plugin(trt_plugins_dir: str, cache_dir: str):
|
|
|
318
332
|
and not os.path.exists(os.path.join(d, "crt", "math_functions.h"))
|
|
319
333
|
]
|
|
320
334
|
|
|
321
|
-
cpp_compiler =
|
|
335
|
+
cpp_compiler = _host_compiler()
|
|
322
336
|
ccbin_flag = ["-ccbin", cpp_compiler] if os.path.exists(f"/usr/bin/{cpp_compiler}") else []
|
|
323
337
|
|
|
324
338
|
cmd = [
|
|
@@ -383,7 +397,7 @@ def build_trt_plugin(trt_plugins_dir: str, cache_dir: str):
|
|
|
383
397
|
nvinfer_link_flag = [os.path.join(d, f)]
|
|
384
398
|
break
|
|
385
399
|
|
|
386
|
-
cpp_compiler =
|
|
400
|
+
cpp_compiler = _host_compiler()
|
|
387
401
|
cmd = [
|
|
388
402
|
cpp_compiler, "-shared", "-o", so_path
|
|
389
403
|
] + obj_files + [
|
|
@@ -150,31 +150,26 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
|
|
|
150
150
|
arg_setup_lines.append("auto ref_in = ctx.GetInput(0);")
|
|
151
151
|
arg_setup_lines.append("auto info = ref_in.GetTensorTypeAndShapeInfo();")
|
|
152
152
|
arg_setup_lines.append("std::vector<int64_t> dim_values = info.GetShape();")
|
|
153
|
+
|
|
154
|
+
# Dims the caller asked to keep dynamic: read them from the tensor that
|
|
155
|
+
# is actually being run, not from the traced shape.
|
|
156
|
+
_dyn_axes = tuple(getattr(manifest, "dynamic_axes", ()) or ())
|
|
153
157
|
|
|
154
158
|
import re
|
|
155
159
|
fn_match = re.search(r'\.entry\s+([a-zA-Z0-9_]+)', manifest.ptx)
|
|
156
160
|
ptx_entry_name = fn_match.group(1) if fn_match else manifest.kernel_name
|
|
157
161
|
|
|
158
|
-
# ---
|
|
159
|
-
grid_strs =
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
expr = str(g.node.expr) if hasattr(g, 'node') else str(g)
|
|
163
|
-
expr = re.sub(r'([a-zA-Z0-9_]+)\*\*([a-zA-Z0-9_]+)', r'std::pow(\1, \2)', expr)
|
|
164
|
-
expr = expr.replace("//", "/")
|
|
165
|
-
expr = re.sub(r'floor\((.*?)\)', r'(\1)', expr)
|
|
166
|
-
|
|
167
|
-
# Replace any s(\d+) symbol dynamically using regex
|
|
168
|
-
expr = re.sub(r'\bs(\d+)\b', r"(int64_t)(dim_values.size() > \1 ? dim_values[\1] : 1)", expr)
|
|
169
|
-
grid_strs.append(expr)
|
|
170
|
-
|
|
162
|
+
# --- GRID EVALUATION, from the manifest's own symbol translation ---
|
|
163
|
+
grid_strs = manifest.grid_exprs(
|
|
164
|
+
lambda i: f"(int64_t)(dim_values.size() > {i} ? dim_values[{i}] : 1)"
|
|
165
|
+
)
|
|
171
166
|
while len(grid_strs) < 3:
|
|
172
167
|
grid_strs.append("1")
|
|
173
168
|
|
|
174
169
|
grid_eval_lines = [
|
|
175
170
|
f"unsigned int grid_x = std::max(1u, (unsigned int)({grid_strs[0]}));",
|
|
176
171
|
f"unsigned int grid_y = std::max(1u, (unsigned int)({grid_strs[1]}));",
|
|
177
|
-
f"unsigned int grid_z = std::max(1u, (unsigned int)({grid_strs[2]}));"
|
|
172
|
+
f"unsigned int grid_z = std::max(1u, (unsigned int)({grid_strs[2]}));",
|
|
178
173
|
]
|
|
179
174
|
|
|
180
175
|
entry_match = re.search(r'\.entry\s+[a-zA-Z0-9_]+\s*\((.*?)\)\s*(?:\.reqntid|\{)', manifest.ptx, re.DOTALL)
|
|
@@ -215,6 +210,23 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
|
|
|
215
210
|
|
|
216
211
|
ptx_ordered_slots.append({"type": c_type, "ident": str(param_num), "is_ptr": is_ptr, "decl": p_decl, "param_num": param_num})
|
|
217
212
|
|
|
213
|
+
# The node's inputs, in the order ONNX Runtime presents them. This is
|
|
214
|
+
# the same list GetInputTypeCount() reports, so indexing into it is
|
|
215
|
+
# authoritative. A running counter over manifest.arguments is not: any
|
|
216
|
+
# drift between the two silently points a scalar read at the wrong
|
|
217
|
+
# slot, and the kernel then masks against a garbage element count.
|
|
218
|
+
inputs_to_node = [
|
|
219
|
+
a for a in manifest.arguments
|
|
220
|
+
if (a.kind in ('input', 'inplace')
|
|
221
|
+
or (a.kind == 'scalar' and not getattr(a, 'is_constexpr', False)))
|
|
222
|
+
]
|
|
223
|
+
|
|
224
|
+
def _node_input_index(a, fallback):
|
|
225
|
+
for k, cand in enumerate(inputs_to_node):
|
|
226
|
+
if cand is a:
|
|
227
|
+
return k
|
|
228
|
+
return fallback
|
|
229
|
+
|
|
218
230
|
arg_setup_lines.append("// Extract inputs/outputs/scalars from ORT context")
|
|
219
231
|
onnx_input_counter = 0
|
|
220
232
|
ort_output_counter = 0
|
|
@@ -227,7 +239,7 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
|
|
|
227
239
|
if _esz == 0.5
|
|
228
240
|
else f"(in_tensor_{slot_idx}.GetTensorTypeAndShapeInfo().GetElementCount() * {int(_esz)})"
|
|
229
241
|
)
|
|
230
|
-
arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({onnx_input_counter});")
|
|
242
|
+
arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({_node_input_index(arg, onnx_input_counter)});")
|
|
231
243
|
arg_setup_lines.append(f"static thread_local const void* arg_ptr_{slot_idx};")
|
|
232
244
|
arg_setup_lines.append(f"const void* raw_ptr_{slot_idx} = (const void*)in_tensor_{slot_idx}.GetTensorData<uint8_t>();")
|
|
233
245
|
arg_setup_lines.append(f"static thread_local void* align_buf_{slot_idx} = nullptr;")
|
|
@@ -242,6 +254,9 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
|
|
|
242
254
|
elif arg.kind == 'output':
|
|
243
255
|
out_dim_exprs = []
|
|
244
256
|
for dim_idx, d in enumerate(arg.shape):
|
|
257
|
+
if dim_idx in _dyn_axes:
|
|
258
|
+
out_dim_exprs.append(f"(int64_t)(dim_values.size() > {dim_idx} ? dim_values[{dim_idx}] : 1)")
|
|
259
|
+
continue
|
|
245
260
|
try:
|
|
246
261
|
d_int = int(d)
|
|
247
262
|
out_dim_exprs.append(str(d_int))
|
|
@@ -257,6 +272,9 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
|
|
|
257
272
|
elif arg.kind == 'inplace':
|
|
258
273
|
out_dim_exprs = []
|
|
259
274
|
for dim_idx, d in enumerate(arg.shape):
|
|
275
|
+
if dim_idx in _dyn_axes:
|
|
276
|
+
out_dim_exprs.append(f"(int64_t)(dim_values.size() > {dim_idx} ? dim_values[{dim_idx}] : 1)")
|
|
277
|
+
continue
|
|
260
278
|
try:
|
|
261
279
|
d_int = int(d)
|
|
262
280
|
out_dim_exprs.append(str(d_int))
|
|
@@ -264,7 +282,7 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
|
|
|
264
282
|
out_dim_exprs.append(f"(int64_t)(dim_values.size() > {dim_idx} ? dim_values[{dim_idx}] : 1)")
|
|
265
283
|
out_shape_str = "{" + ", ".join(out_dim_exprs) + "}" if out_dim_exprs else "dim_values"
|
|
266
284
|
|
|
267
|
-
arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({onnx_input_counter});")
|
|
285
|
+
arg_setup_lines.append(f"auto in_tensor_{slot_idx} = ctx.GetInput({_node_input_index(arg, onnx_input_counter)});")
|
|
268
286
|
arg_setup_lines.append(f"std::vector<int64_t> out_dims_{ort_output_counter} = {out_shape_str};")
|
|
269
287
|
arg_setup_lines.append(f"auto out_tensor_{slot_idx} = ctx.GetOutput({ort_output_counter}, out_dims_{ort_output_counter}.data(), out_dims_{ort_output_counter}.size());")
|
|
270
288
|
_esz = element_size_bytes(arg.dtype)
|
|
@@ -282,16 +300,33 @@ struct {op_name} : Ort::CustomOpBase<{op_name}, {kernel_name}> {{
|
|
|
282
300
|
onnx_input_counter += 1
|
|
283
301
|
ort_output_counter += 1
|
|
284
302
|
elif arg.kind == 'scalar':
|
|
285
|
-
|
|
286
|
-
|
|
303
|
+
expr = getattr(arg, 'cxx_expr', '') or ''
|
|
304
|
+
if 'dim_values' in expr:
|
|
305
|
+
# Shape-derived (n = x.numel()): compute it from the tensor
|
|
306
|
+
# being run. The exported graph folds this scalar into a
|
|
307
|
+
# constant -- it is not a graph input -- so reading the slot
|
|
308
|
+
# would hand the kernel the element count captured while
|
|
309
|
+
# tracing, and it would then process a fraction of the
|
|
310
|
+
# tensor and leave the rest untouched.
|
|
311
|
+
arg_setup_lines.append(f"double scalar_val_{slot_idx} = (double)({expr});")
|
|
312
|
+
onnx_input_counter += 1
|
|
313
|
+
elif not getattr(arg, 'is_constexpr', False):
|
|
314
|
+
arg_setup_lines.append(f"auto scalar_tensor_{slot_idx} = ctx.GetInput({_node_input_index(arg, onnx_input_counter)});")
|
|
287
315
|
arg_setup_lines.append(f"double scalar_val_{slot_idx} = 0;")
|
|
288
316
|
arg_setup_lines.append(f"if (scalar_tensor_{slot_idx}.GetTensorTypeAndShapeInfo().GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) {{ scalar_val_{slot_idx} = (double)(scalar_tensor_{slot_idx}.GetTensorData<int64_t>()[0]); }}")
|
|
289
317
|
arg_setup_lines.append(f"else if (scalar_tensor_{slot_idx}.GetTensorTypeAndShapeInfo().GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32) {{ scalar_val_{slot_idx} = (double)(scalar_tensor_{slot_idx}.GetTensorData<int32_t>()[0]); }}")
|
|
290
318
|
arg_setup_lines.append(f"else {{ scalar_val_{slot_idx} = (double)(scalar_tensor_{slot_idx}.GetTensorData<float>()[0]); }}")
|
|
291
319
|
onnx_input_counter += 1
|
|
292
320
|
else:
|
|
293
|
-
|
|
294
|
-
|
|
321
|
+
# Prefer the manifest's translated expression: for a scalar
|
|
322
|
+
# derived from shapes (n = x.numel()) that is a runtime read
|
|
323
|
+
# of the input dims; otherwise it is the baked constant.
|
|
324
|
+
expr = getattr(arg, 'cxx_expr', '') or None
|
|
325
|
+
if expr:
|
|
326
|
+
arg_setup_lines.append(f"double scalar_val_{slot_idx} = (double)({expr});")
|
|
327
|
+
else:
|
|
328
|
+
s_val = arg.value if getattr(arg, 'value', None) is not None else 0.0
|
|
329
|
+
arg_setup_lines.append(f"double scalar_val_{slot_idx} = {s_val};")
|
|
295
330
|
|
|
296
331
|
arg_setup_lines.append("static thread_local const void* null_ptr_arg = nullptr;")
|
|
297
332
|
arg_setup_lines.append("static thread_local std::vector<void*> kp;")
|