kernel-lens 1.2.0__tar.gz → 1.2.1__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.
Files changed (58) hide show
  1. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/PKG-INFO +105 -29
  2. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/README.md +88 -26
  3. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/__init__.py +1 -1
  4. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/builder.py +17 -3
  5. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/runtime/engine.py +95 -8
  6. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/PKG-INFO +105 -29
  7. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/requires.txt +3 -0
  8. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/pyproject.toml +27 -3
  9. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/LICENSE +0 -0
  10. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/__init__.py +0 -0
  11. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/ort_gen.py +0 -0
  12. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/trt_gen.py +0 -0
  13. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/__init__.py +0 -0
  14. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/ast_analyzer.py +0 -0
  15. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/core.py +0 -0
  16. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/interaction.py +0 -0
  17. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/manifest.py +0 -0
  18. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/onnx_exporter.py +0 -0
  19. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/tracer.py +0 -0
  20. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/config.py +0 -0
  21. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInfer.h +0 -0
  22. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferImpl.h +0 -0
  23. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferLegacyDims.h +0 -0
  24. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferPlugin.h +0 -0
  25. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferPluginBase.h +0 -0
  26. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferPluginUtils.h +0 -0
  27. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntime.h +0 -0
  28. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntimeBase.h +0 -0
  29. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntimeCommon.h +0 -0
  30. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntimePlugin.h +0 -0
  31. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferVersion.h +0 -0
  32. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/runtime/__init__.py +0 -0
  33. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/utils/__init__.py +0 -0
  34. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/utils/deployment.py +0 -0
  35. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/utils/env_check.py +0 -0
  36. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/SOURCES.txt +0 -0
  37. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/dependency_links.txt +0 -0
  38. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/top_level.txt +0 -0
  39. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/setup.cfg +0 -0
  40. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_alignment_guard.py +0 -0
  41. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_architectural_fixes.py +0 -0
  42. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_conformance.py +0 -0
  43. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_edge_cases.py +0 -0
  44. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_fp8_gemm.py +0 -0
  45. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_fused_cross_entropy.py +0 -0
  46. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_fused_residual_rms.py +0 -0
  47. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_inplace.py +0 -0
  48. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_inplace_dtypes.py +0 -0
  49. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_int4_dequant_gemm.py +0 -0
  50. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_mamba_selective_scan.py +0 -0
  51. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_research_kernels.py +0 -0
  52. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_standalone.py +0 -0
  53. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_standalone_launch.py +0 -0
  54. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_standalone_ptx.py +0 -0
  55. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_tricky_triton_kernels.py +0 -0
  56. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_triton_vs_torch_compile.py +0 -0
  57. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_trt.py +0 -0
  58. {kernel_lens-1.2.0 → kernel_lens-1.2.1}/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.2.0
4
- Summary: A production-grade, multi-backend compiler for Triton kernels (ONNX & TensorRT)
5
- Author: Compiler Engineer
3
+ Version: 1.2.1
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
@@ -113,11 +127,16 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
113
127
  dynamic scalars, strided and dynamic inputs, loops, conditionals, nested
114
128
  helpers, in-place mutation, 3-D grids, autotune, and fp16/bf16.
115
129
 
116
- > **Honest scope.** KernelLens is research-grade. The documentation clearly
117
- > separates what is **implemented and measured** from what is **planned**. Two
118
- > TensorRT paths (in-place aliasing, INT8 dequantisation) are known to be
119
- > incorrect and are reported rather than hidden. See the
120
- > [verification page](https://just1truc.github.io/kernel_lens/verification/).
130
+ > **Honest scope.** KernelLens is research-grade, and the results are reported
131
+ > per configuration rather than blended. The composition is bit-exact against
132
+ > native Triton (`0.00e+00`) across the research kernels, the fourteen
133
+ > conformance patterns and the quantized GEMMs, on both backends; two
134
+ > multi-operator fusions land just above zero and are reported as such. Keeping
135
+ > the kernel pays when the baseline leaves the graph unfused and breaks even when
136
+ > it already fuses. The TensorRT-LLM backend is design-only (see
137
+ > [`docs/assets/data/a100/trtllm_backend_design.md`](docs/assets/data/a100/trtllm_backend_design.md)),
138
+ > and the quantized GEMMs are validated on ONNX Runtime only. Full methodology
139
+ > and every raw JSON are under [`docs/assets/data/`](docs/assets/data/).
121
140
 
122
141
  ## Installation
123
142
 
@@ -129,7 +148,16 @@ pip install kernel-lens[all] # both
129
148
  ```
130
149
 
131
150
  **Requirements:** Linux x86-64, Python ≥ 3.10, CUDA ≥ 12.0, PyTorch ≥ 2.0,
132
- Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`.
151
+ Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
152
+ generator emits `IPluginV2DynamicExt`).
153
+
154
+ **Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
155
+ `.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
156
+ compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
157
+ specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
158
+ `TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
159
+ above 2 GB from disk rather than from a memory buffer (opt-in; see the
160
+ [changelog](CHANGELOG.md)).
133
161
 
134
162
  ## Quickstart
135
163
 
@@ -172,22 +200,52 @@ out_trt = compiled.run((x,), backend="tensorrt")
172
200
 
173
201
  ## Results at a glance
174
202
 
175
- Measured on an **RTX 5060 Ti (Blackwell, `sm_120a`)**, CUDA 12.8, PyTorch 2.14,
176
- Triton 3.8, ONNX Runtime 1.30, TensorRT 10.16. Full tables and methodology are on
177
- the [results page](https://just1truc.github.io/kernel_lens/#results).
178
-
179
- - **Fold + keep, and it pays**: on a two-layer block whose attention is a custom
180
- squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
181
- for the folded-only `torch.compile` formulation and **28.04 ms** for eager
182
- folding — **1.99×** and **3.68×** faster — and **1.08×** faster than the
183
- unfused Triton-kernel path. The engine folds the standard layers; the kernel
184
- keeps the operations folding cannot express.
185
- - **Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
186
- PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and
187
- the quantized GEMMs, on both backends.
188
- - **Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
189
- and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
190
- 7.85/7.92 ms).
203
+ Two architectures: an **RTX 5060 Ti** (Blackwell, `sm_120a`) for the folding,
204
+ single-operator, end-to-end and quantized results, and an **A100-SXM4-80GB**
205
+ (Ampere, `sm_80`) for the conformance, parity and real-model results. Results are
206
+ reported per configuration and never blended. Full methodology, complete tables
207
+ and every raw JSON are on the
208
+ [results page](https://just1truc.github.io/kernel_lens/#results) and under
209
+ [`docs/assets/data/`](docs/assets/data/).
210
+
211
+ **The code you don't write.** The deployment is **nine lines of Python**; the C++
212
+ it replaces is **1,137 lines across six files** for one fused RMSNorm kernel on
213
+ both backends, measured with `tests/measure_codegen.py`.
214
+
215
+ **Real pretrained models.** Every decoder `LayerNorm` of a real OPT checkpoint
216
+ (125M → 6.7B) replaced by the Triton RMSNorm kernel, graded against
217
+ `torch.compile` — the baseline that *captures the same Triton kernels* rather
218
+ than reimplementing them. The result splits by precision:
219
+
220
+ - **fp32** — `torch.compile` barely improves on eager (1.02–1.12×) because it
221
+ leaves the operators around the kernel unfused. KernelLens fuses them:
222
+ **3.18–5.45× faster** under TensorRT and **2.27–4.62×** under ONNX Runtime,
223
+ and the margin grows with model size.
224
+ - **fp16** — `torch.compile` optimises well there, so we only match it: TensorRT
225
+ edges it by **1.06–1.22×**, while ONNX Runtime is consistently **slower
226
+ (0.67–0.84×)** — per-node dispatch dominates when the graph already fuses.
227
+ - **Accuracy tracks precision, not speed.** fp32 holds the same greedy token on
228
+ **≥ 98.8%** of positions; fp16 under TensorRT falls to 93.4–97.9%. OPT-350M in
229
+ fp16 under ONNX Runtime is bit-exact.
230
+
231
+ **Fold + keep, and it pays.** On a two-layer block whose attention is a custom
232
+ squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
233
+ for the folded-only `torch.compile` formulation and **28.04 ms** for eager
234
+ folding — **1.99×** and **3.68×** faster. Against the unfused kernel path it is
235
+ **1.08×** faster than eager Triton (8.24 ms) and level with `torch.compile`
236
+ (7.55 ms). The engine folds the standard layers; the kernel keeps the operations
237
+ folding cannot express.
238
+
239
+ **Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
240
+ PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and the
241
+ quantized GEMMs, on both backends — and across all **fourteen** module patterns
242
+ of `tests/test_conformance.py`. Two multi-operator fusions are near-exact rather
243
+ than exact (`≤ 9.5e-07` residual+RMSNorm, `4.8e-07` Mamba's scan) and are
244
+ reported as such.
245
+
246
+ **Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
247
+ and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
248
+ 7.85/7.92 ms).
191
249
 
192
250
  ## Documentation
193
251
 
@@ -195,9 +253,11 @@ the [results page](https://just1truc.github.io/kernel_lens/#results).
195
253
  | :--- | :--- |
196
254
  | [Problem](https://just1truc.github.io/kernel_lens/#problem) | The five gaps a naive export hits |
197
255
  | [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, latency, end-to-end |
256
+ | [Results](https://just1truc.github.io/kernel_lens/#results) | Folding + kernels, parity, conformance, single-operator and end-to-end latency, **real pretrained models** |
257
+ | [Why an engine](https://just1truc.github.io/kernel_lens/#why) | `torch.compile` vs TensorRT, both directions |
199
258
  | [Install & API](https://just1truc.github.io/kernel_lens/#install) | Installation and API reference |
200
259
  | [Status](https://just1truc.github.io/kernel_lens/#status) | Roadmap |
260
+ | [FAQ](https://just1truc.github.io/kernel_lens/#faq) | Fair objections |
201
261
  | [Paper](https://just1truc.github.io/kernel_lens/#paper) | Abstract and citation |
202
262
 
203
263
  ## Reproduce the results
@@ -212,14 +272,30 @@ python tests/test_research_kernels.py # numerical parity
212
272
  python tests/benchmark_research_kernels.py # single-operator latency
213
273
  python tests/benchmark_e2e_llama.py # end-to-end decoder
214
274
  python tests/test_fp8_gemm.py # FP8 Tensor Core path
275
+ python tests/measure_codegen.py # generated C++ volume, per file
276
+ python tests/benchmark_real_model.py --model facebook/opt-1.3b # real OPT, both precisions
215
277
  ```
216
278
 
279
+ The real-model matrix (OPT-125M → 6.7B, fp32 and fp16) runs unattended on a fresh
280
+ CUDA host: `scripts/remote_setup.sh` provisions nvcc, TensorRT, `onnxruntime-gpu`
281
+ and `transformers` and installs the package; `scripts/remote_bench_all.sh` then
282
+ sweeps the matrix into `results/` and writes `summary.json`, recording failures
283
+ instead of aborting. See
284
+ [`docs/assets/data/a100/runpod_runbook.md`](docs/assets/data/a100/runpod_runbook.md) for the
285
+ exact invocation and the environment traps (CUDA 12.8's `nvcc` does not build
286
+ against a modern libstdc++; CUDA 13.3's does).
287
+
217
288
  Raw measurement JSON is committed under [`docs/assets/data/`](docs/assets/data/).
218
289
 
219
290
  ## Technical report
220
291
 
221
- The LaTeX source is [`architecture_report.tex`](architecture_report.tex); the TMLR
222
- variant is generated from it with `python3 scripts/make_tmlr.py`.
292
+ The paper source is [`kernel_lens_fresh.tex`](kernel_lens_fresh.tex). Two derived
293
+ variants are generated from it and must not be edited by hand:
294
+ `scripts/make_mlsys.py` produces the 2-column MLSys form
295
+ (`kernel_lens_mlsys.tex`) and `scripts/make_tmlr_fresh.py` produces the TMLR
296
+ forms — `kernel_lens_tmlr_fresh.tex` (named preprint) and
297
+ `kernel_lens_tmlr_submission.tex` (anonymized; **do not publish it alongside the
298
+ named preprint**).
223
299
 
224
300
  ```bibtex
225
301
  @article{duc2026kernellens,
@@ -227,7 +303,7 @@ variant is generated from it with `python3 scripts/make_tmlr.py`.
227
303
  author = {Duc, Justin},
228
304
  institution = {Galiad Research},
229
305
  year = {2026},
230
- version = {1.1.9},
306
+ version = {1.2.1},
231
307
  doi = {10.5281/zenodo.23194523},
232
308
  url = {https://doi.org/10.5281/zenodo.23194523}
233
309
  }
@@ -63,11 +63,16 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
63
63
  dynamic scalars, strided and dynamic inputs, loops, conditionals, nested
64
64
  helpers, in-place mutation, 3-D grids, autotune, and fp16/bf16.
65
65
 
66
- > **Honest scope.** KernelLens is research-grade. The documentation clearly
67
- > separates what is **implemented and measured** from what is **planned**. Two
68
- > TensorRT paths (in-place aliasing, INT8 dequantisation) are known to be
69
- > incorrect and are reported rather than hidden. See the
70
- > [verification page](https://just1truc.github.io/kernel_lens/verification/).
66
+ > **Honest scope.** KernelLens is research-grade, and the results are reported
67
+ > per configuration rather than blended. The composition is bit-exact against
68
+ > native Triton (`0.00e+00`) across the research kernels, the fourteen
69
+ > conformance patterns and the quantized GEMMs, on both backends; two
70
+ > multi-operator fusions land just above zero and are reported as such. Keeping
71
+ > the kernel pays when the baseline leaves the graph unfused and breaks even when
72
+ > it already fuses. The TensorRT-LLM backend is design-only (see
73
+ > [`docs/assets/data/a100/trtllm_backend_design.md`](docs/assets/data/a100/trtllm_backend_design.md)),
74
+ > and the quantized GEMMs are validated on ONNX Runtime only. Full methodology
75
+ > and every raw JSON are under [`docs/assets/data/`](docs/assets/data/).
71
76
 
72
77
  ## Installation
73
78
 
@@ -79,7 +84,16 @@ pip install kernel-lens[all] # both
79
84
  ```
80
85
 
81
86
  **Requirements:** Linux x86-64, Python ≥ 3.10, CUDA ≥ 12.0, PyTorch ≥ 2.0,
82
- Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`.
87
+ Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
88
+ generator emits `IPluginV2DynamicExt`).
89
+
90
+ **Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
91
+ `.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
92
+ compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
93
+ specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
94
+ `TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
95
+ above 2 GB from disk rather than from a memory buffer (opt-in; see the
96
+ [changelog](CHANGELOG.md)).
83
97
 
84
98
  ## Quickstart
85
99
 
@@ -122,22 +136,52 @@ out_trt = compiled.run((x,), backend="tensorrt")
122
136
 
123
137
  ## Results at a glance
124
138
 
125
- Measured on an **RTX 5060 Ti (Blackwell, `sm_120a`)**, CUDA 12.8, PyTorch 2.14,
126
- Triton 3.8, ONNX Runtime 1.30, TensorRT 10.16. Full tables and methodology are on
127
- the [results page](https://just1truc.github.io/kernel_lens/#results).
128
-
129
- - **Fold + keep, and it pays**: on a two-layer block whose attention is a custom
130
- squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
131
- for the folded-only `torch.compile` formulation and **28.04 ms** for eager
132
- folding — **1.99×** and **3.68×** faster — and **1.08×** faster than the
133
- unfused Triton-kernel path. The engine folds the standard layers; the kernel
134
- keeps the operations folding cannot express.
135
- - **Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
136
- PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and
137
- the quantized GEMMs, on both backends.
138
- - **Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
139
- and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
140
- 7.85/7.92 ms).
139
+ Two architectures: an **RTX 5060 Ti** (Blackwell, `sm_120a`) for the folding,
140
+ single-operator, end-to-end and quantized results, and an **A100-SXM4-80GB**
141
+ (Ampere, `sm_80`) for the conformance, parity and real-model results. Results are
142
+ reported per configuration and never blended. Full methodology, complete tables
143
+ and every raw JSON are on the
144
+ [results page](https://just1truc.github.io/kernel_lens/#results) and under
145
+ [`docs/assets/data/`](docs/assets/data/).
146
+
147
+ **The code you don't write.** The deployment is **nine lines of Python**; the C++
148
+ it replaces is **1,137 lines across six files** for one fused RMSNorm kernel on
149
+ both backends, measured with `tests/measure_codegen.py`.
150
+
151
+ **Real pretrained models.** Every decoder `LayerNorm` of a real OPT checkpoint
152
+ (125M → 6.7B) replaced by the Triton RMSNorm kernel, graded against
153
+ `torch.compile` — the baseline that *captures the same Triton kernels* rather
154
+ than reimplementing them. The result splits by precision:
155
+
156
+ - **fp32** — `torch.compile` barely improves on eager (1.02–1.12×) because it
157
+ leaves the operators around the kernel unfused. KernelLens fuses them:
158
+ **3.18–5.45× faster** under TensorRT and **2.27–4.62×** under ONNX Runtime,
159
+ and the margin grows with model size.
160
+ - **fp16** — `torch.compile` optimises well there, so we only match it: TensorRT
161
+ edges it by **1.06–1.22×**, while ONNX Runtime is consistently **slower
162
+ (0.67–0.84×)** — per-node dispatch dominates when the graph already fuses.
163
+ - **Accuracy tracks precision, not speed.** fp32 holds the same greedy token on
164
+ **≥ 98.8%** of positions; fp16 under TensorRT falls to 93.4–97.9%. OPT-350M in
165
+ fp16 under ONNX Runtime is bit-exact.
166
+
167
+ **Fold + keep, and it pays.** On a two-layer block whose attention is a custom
168
+ squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
169
+ for the folded-only `torch.compile` formulation and **28.04 ms** for eager
170
+ folding — **1.99×** and **3.68×** faster. Against the unfused kernel path it is
171
+ **1.08×** faster than eager Triton (8.24 ms) and level with `torch.compile`
172
+ (7.55 ms). The engine folds the standard layers; the kernel keeps the operations
173
+ folding cannot express.
174
+
175
+ **Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
176
+ PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and the
177
+ quantized GEMMs, on both backends — and across all **fourteen** module patterns
178
+ of `tests/test_conformance.py`. Two multi-operator fusions are near-exact rather
179
+ than exact (`≤ 9.5e-07` residual+RMSNorm, `4.8e-07` Mamba's scan) and are
180
+ reported as such.
181
+
182
+ **Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
183
+ and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
184
+ 7.85/7.92 ms).
141
185
 
142
186
  ## Documentation
143
187
 
@@ -145,9 +189,11 @@ the [results page](https://just1truc.github.io/kernel_lens/#results).
145
189
  | :--- | :--- |
146
190
  | [Problem](https://just1truc.github.io/kernel_lens/#problem) | The five gaps a naive export hits |
147
191
  | [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, latency, end-to-end |
192
+ | [Results](https://just1truc.github.io/kernel_lens/#results) | Folding + kernels, parity, conformance, single-operator and end-to-end latency, **real pretrained models** |
193
+ | [Why an engine](https://just1truc.github.io/kernel_lens/#why) | `torch.compile` vs TensorRT, both directions |
149
194
  | [Install & API](https://just1truc.github.io/kernel_lens/#install) | Installation and API reference |
150
195
  | [Status](https://just1truc.github.io/kernel_lens/#status) | Roadmap |
196
+ | [FAQ](https://just1truc.github.io/kernel_lens/#faq) | Fair objections |
151
197
  | [Paper](https://just1truc.github.io/kernel_lens/#paper) | Abstract and citation |
152
198
 
153
199
  ## Reproduce the results
@@ -162,14 +208,30 @@ python tests/test_research_kernels.py # numerical parity
162
208
  python tests/benchmark_research_kernels.py # single-operator latency
163
209
  python tests/benchmark_e2e_llama.py # end-to-end decoder
164
210
  python tests/test_fp8_gemm.py # FP8 Tensor Core path
211
+ python tests/measure_codegen.py # generated C++ volume, per file
212
+ python tests/benchmark_real_model.py --model facebook/opt-1.3b # real OPT, both precisions
165
213
  ```
166
214
 
215
+ The real-model matrix (OPT-125M → 6.7B, fp32 and fp16) runs unattended on a fresh
216
+ CUDA host: `scripts/remote_setup.sh` provisions nvcc, TensorRT, `onnxruntime-gpu`
217
+ and `transformers` and installs the package; `scripts/remote_bench_all.sh` then
218
+ sweeps the matrix into `results/` and writes `summary.json`, recording failures
219
+ instead of aborting. See
220
+ [`docs/assets/data/a100/runpod_runbook.md`](docs/assets/data/a100/runpod_runbook.md) for the
221
+ exact invocation and the environment traps (CUDA 12.8's `nvcc` does not build
222
+ against a modern libstdc++; CUDA 13.3's does).
223
+
167
224
  Raw measurement JSON is committed under [`docs/assets/data/`](docs/assets/data/).
168
225
 
169
226
  ## Technical report
170
227
 
171
- The LaTeX source is [`architecture_report.tex`](architecture_report.tex); the TMLR
172
- variant is generated from it with `python3 scripts/make_tmlr.py`.
228
+ The paper source is [`kernel_lens_fresh.tex`](kernel_lens_fresh.tex). Two derived
229
+ variants are generated from it and must not be edited by hand:
230
+ `scripts/make_mlsys.py` produces the 2-column MLSys form
231
+ (`kernel_lens_mlsys.tex`) and `scripts/make_tmlr_fresh.py` produces the TMLR
232
+ forms — `kernel_lens_tmlr_fresh.tex` (named preprint) and
233
+ `kernel_lens_tmlr_submission.tex` (anonymized; **do not publish it alongside the
234
+ named preprint**).
173
235
 
174
236
  ```bibtex
175
237
  @article{duc2026kernellens,
@@ -177,7 +239,7 @@ variant is generated from it with `python3 scripts/make_tmlr.py`.
177
239
  author = {Duc, Justin},
178
240
  institution = {Galiad Research},
179
241
  year = {2026},
180
- version = {1.1.9},
242
+ version = {1.2.1},
181
243
  doi = {10.5281/zenodo.23194523},
182
244
  url = {https://doi.org/10.5281/zenodo.23194523}
183
245
  }
@@ -7,7 +7,7 @@ from .config import set_verbose, is_verbose
7
7
 
8
8
 
9
9
 
10
- __version__ = "1.2.0"
10
+ __version__ = "1.2.1"
11
11
 
12
12
  __all__ = [
13
13
  "compile",
@@ -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 = "g++-13" if os.path.exists("/usr/bin/g++-13") else "g++"
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 = "g++-13" if os.path.exists("/usr/bin/g++-13") else "g++"
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 = "g++-13" if os.path.exists("/usr/bin/g++-13") else "g++"
400
+ cpp_compiler = _host_compiler()
387
401
  cmd = [
388
402
  cpp_compiler, "-shared", "-o", so_path
389
403
  ] + obj_files + [
@@ -31,6 +31,63 @@ def _graph_has_int8_custom_inputs(onnx_path: str) -> bool:
31
31
  return False
32
32
  return False
33
33
 
34
+ def _make_stream_reader(trt, path):
35
+ """File-backed ``IStreamReaderV2``.
36
+
37
+ Engines above roughly 2 GB fail to deserialize from a memory buffer; the
38
+ stream-reader path reads them straight from disk instead. TensorRT fills a
39
+ caller-provided capsule, so the bytes have to be memmove'd into it.
40
+ """
41
+ import ctypes
42
+
43
+ _get_ptr = ctypes.pythonapi.PyCapsule_GetPointer
44
+ _get_ptr.restype = ctypes.c_void_p
45
+ _get_ptr.argtypes = [ctypes.py_object, ctypes.c_char_p]
46
+
47
+ _whence = {trt.SeekPosition.SET: os.SEEK_SET,
48
+ trt.SeekPosition.CUR: os.SEEK_CUR,
49
+ trt.SeekPosition.END: os.SEEK_END}
50
+
51
+ class _FileStreamReader(trt.IStreamReaderV2):
52
+ def __init__(self):
53
+ super().__init__()
54
+ self._f = open(path, "rb")
55
+
56
+ def read(self, num_bytes, destination):
57
+ # The binding dispatches (num_bytes, destination_pointer) — the
58
+ # reverse of the order in the published docstring, which is itself
59
+ # swapped with seek's. Verified empirically against TensorRT 10.16.
60
+ data = self._f.read(int(num_bytes))
61
+ if not data:
62
+ return 0
63
+ ctypes.memmove(int(destination), data, len(data))
64
+ return len(data)
65
+
66
+ def seek(self, offset, where):
67
+ self._f.seek(int(offset), _whence.get(where, os.SEEK_SET))
68
+ return True
69
+
70
+ return _FileStreamReader()
71
+
72
+
73
+ def _deserialize_engine(trt, runtime, engine_path):
74
+ # The stream-reader path is experimental: TensorRT's Python binding
75
+ # dispatches read(num_bytes, destination) — the reverse of its documented
76
+ # order — and writing through that pointer is not yet reliable. Engines
77
+ # below 2 GB take the memory path regardless; above it, opt in explicitly.
78
+ if (os.environ.get("KERNEL_LENS_STREAM_READER") == "1"
79
+ and os.path.getsize(engine_path) > (1 << 31)
80
+ and hasattr(trt, "IStreamReaderV2")):
81
+ try:
82
+ eng = runtime.deserialize_cuda_engine(_make_stream_reader(trt, engine_path))
83
+ if eng is not None:
84
+ return eng
85
+ except Exception:
86
+ pass
87
+ with open(engine_path, "rb") as f:
88
+ return runtime.deserialize_cuda_engine(f.read())
89
+
90
+
34
91
  class CompiledModel:
35
92
  def __init__(self, cache_dir: str, model_name: str, backends: list[str], output_shapes: list = None, output_strides: list = None, output_dtypes: list = None, saved_args: list = None):
36
93
  """
@@ -306,16 +363,46 @@ class CompiledModel:
306
363
  else:
307
364
  network = builder.create_network()
308
365
  parser = trt.OnnxParser(network, TRT_LOGGER)
309
-
310
- with open(onnx_path, 'rb') as model:
311
- if not parser.parse(model.read()):
312
- for error in range(parser.num_errors):
313
- print(parser.get_error(error))
314
- raise RuntimeError("Failed to parse ONNX graph for TensorRT.")
366
+
367
+ # Parse from the file path rather than a bytes buffer. Graphs
368
+ # above the 2 GB protobuf limit spill their weights to external
369
+ # files beside the model, and the parser can only resolve those
370
+ # relative locations when it knows which directory that is.
371
+ if hasattr(parser, "parse_from_file"):
372
+ ok = parser.parse_from_file(onnx_path)
373
+ else:
374
+ with open(onnx_path, 'rb') as model:
375
+ ok = parser.parse(model.read())
376
+ if not ok:
377
+ for error in range(parser.num_errors):
378
+ print(parser.get_error(error))
379
+ raise RuntimeError("Failed to parse ONNX graph for TensorRT.")
315
380
 
316
381
  config = builder.create_builder_config()
317
382
  if hasattr(config, "set_memory_pool_limit"):
318
383
  config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
384
+
385
+ # TensorRT keeps plugin layers in FP32 unless the matching
386
+ # precision flag is enabled. Without it the builder inserts a
387
+ # cast on the plugin's inputs, and the plugin's dtype guard
388
+ # (which mirrors the kernel's signature) then rejects them, so
389
+ # the engine build fails with "could not find any supported
390
+ # formats consistent with input/output data types".
391
+ try:
392
+ import onnx as _onnx
393
+ _g = _onnx.load(onnx_path, load_external_data=False).graph
394
+ # TensorProto: 10 = FLOAT16, 16 = BFLOAT16. An LLM's graph
395
+ # *inputs* are token ids (int64), so the precision has to be
396
+ # read from the initializers and the internal tensors too.
397
+ _elem = {t.type.tensor_type.elem_type
398
+ for t in list(_g.input) + list(_g.output) + list(_g.value_info)}
399
+ _elem |= {t.data_type for t in _g.initializer}
400
+ if 10 in _elem and hasattr(trt.BuilderFlag, "FP16"):
401
+ config.set_flag(trt.BuilderFlag.FP16)
402
+ if 16 in _elem and hasattr(trt.BuilderFlag, "BF16"):
403
+ config.set_flag(trt.BuilderFlag.BF16)
404
+ except Exception:
405
+ pass
319
406
 
320
407
  profile = builder.create_optimization_profile()
321
408
  has_dynamic = False
@@ -358,8 +445,8 @@ class CompiledModel:
358
445
  # ---------------------------------------------------
359
446
 
360
447
  # On lit le SSD et on désérialise UNE SEULE FOIS
361
- with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
362
- self._trt_engine = runtime.deserialize_cuda_engine(f.read())
448
+ with trt.Runtime(TRT_LOGGER) as runtime:
449
+ self._trt_engine = _deserialize_engine(trt, runtime, engine_path)
363
450
 
364
451
  self._trt_context = self._trt_engine.create_execution_context()
365
452
  self._output_tensors = {} # (name, shape, dtype) -> cached CUDA tensor
@@ -1,8 +1,8 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: kernel-lens
3
- Version: 1.2.0
4
- Summary: A production-grade, multi-backend compiler for Triton kernels (ONNX & TensorRT)
5
- Author: Compiler Engineer
3
+ Version: 1.2.1
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
@@ -113,11 +127,16 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
113
127
  dynamic scalars, strided and dynamic inputs, loops, conditionals, nested
114
128
  helpers, in-place mutation, 3-D grids, autotune, and fp16/bf16.
115
129
 
116
- > **Honest scope.** KernelLens is research-grade. The documentation clearly
117
- > separates what is **implemented and measured** from what is **planned**. Two
118
- > TensorRT paths (in-place aliasing, INT8 dequantisation) are known to be
119
- > incorrect and are reported rather than hidden. See the
120
- > [verification page](https://just1truc.github.io/kernel_lens/verification/).
130
+ > **Honest scope.** KernelLens is research-grade, and the results are reported
131
+ > per configuration rather than blended. The composition is bit-exact against
132
+ > native Triton (`0.00e+00`) across the research kernels, the fourteen
133
+ > conformance patterns and the quantized GEMMs, on both backends; two
134
+ > multi-operator fusions land just above zero and are reported as such. Keeping
135
+ > the kernel pays when the baseline leaves the graph unfused and breaks even when
136
+ > it already fuses. The TensorRT-LLM backend is design-only (see
137
+ > [`docs/assets/data/a100/trtllm_backend_design.md`](docs/assets/data/a100/trtllm_backend_design.md)),
138
+ > and the quantized GEMMs are validated on ONNX Runtime only. Full methodology
139
+ > and every raw JSON are under [`docs/assets/data/`](docs/assets/data/).
121
140
 
122
141
  ## Installation
123
142
 
@@ -129,7 +148,16 @@ pip install kernel-lens[all] # both
129
148
  ```
130
149
 
131
150
  **Requirements:** Linux x86-64, Python ≥ 3.10, CUDA ≥ 12.0, PyTorch ≥ 2.0,
132
- Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`.
151
+ Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
152
+ generator emits `IPluginV2DynamicExt`).
153
+
154
+ **Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
155
+ `.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
156
+ compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
157
+ specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
158
+ `TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
159
+ above 2 GB from disk rather than from a memory buffer (opt-in; see the
160
+ [changelog](CHANGELOG.md)).
133
161
 
134
162
  ## Quickstart
135
163
 
@@ -172,22 +200,52 @@ out_trt = compiled.run((x,), backend="tensorrt")
172
200
 
173
201
  ## Results at a glance
174
202
 
175
- Measured on an **RTX 5060 Ti (Blackwell, `sm_120a`)**, CUDA 12.8, PyTorch 2.14,
176
- Triton 3.8, ONNX Runtime 1.30, TensorRT 10.16. Full tables and methodology are on
177
- the [results page](https://just1truc.github.io/kernel_lens/#results).
178
-
179
- - **Fold + keep, and it pays**: on a two-layer block whose attention is a custom
180
- squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
181
- for the folded-only `torch.compile` formulation and **28.04 ms** for eager
182
- folding — **1.99×** and **3.68×** faster — and **1.08×** faster than the
183
- unfused Triton-kernel path. The engine folds the standard layers; the kernel
184
- keeps the operations folding cannot express.
185
- - **Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
186
- PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and
187
- the quantized GEMMs, on both backends.
188
- - **Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
189
- and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
190
- 7.85/7.92 ms).
203
+ Two architectures: an **RTX 5060 Ti** (Blackwell, `sm_120a`) for the folding,
204
+ single-operator, end-to-end and quantized results, and an **A100-SXM4-80GB**
205
+ (Ampere, `sm_80`) for the conformance, parity and real-model results. Results are
206
+ reported per configuration and never blended. Full methodology, complete tables
207
+ and every raw JSON are on the
208
+ [results page](https://just1truc.github.io/kernel_lens/#results) and under
209
+ [`docs/assets/data/`](docs/assets/data/).
210
+
211
+ **The code you don't write.** The deployment is **nine lines of Python**; the C++
212
+ it replaces is **1,137 lines across six files** for one fused RMSNorm kernel on
213
+ both backends, measured with `tests/measure_codegen.py`.
214
+
215
+ **Real pretrained models.** Every decoder `LayerNorm` of a real OPT checkpoint
216
+ (125M → 6.7B) replaced by the Triton RMSNorm kernel, graded against
217
+ `torch.compile` — the baseline that *captures the same Triton kernels* rather
218
+ than reimplementing them. The result splits by precision:
219
+
220
+ - **fp32** — `torch.compile` barely improves on eager (1.02–1.12×) because it
221
+ leaves the operators around the kernel unfused. KernelLens fuses them:
222
+ **3.18–5.45× faster** under TensorRT and **2.27–4.62×** under ONNX Runtime,
223
+ and the margin grows with model size.
224
+ - **fp16** — `torch.compile` optimises well there, so we only match it: TensorRT
225
+ edges it by **1.06–1.22×**, while ONNX Runtime is consistently **slower
226
+ (0.67–0.84×)** — per-node dispatch dominates when the graph already fuses.
227
+ - **Accuracy tracks precision, not speed.** fp32 holds the same greedy token on
228
+ **≥ 98.8%** of positions; fp16 under TensorRT falls to 93.4–97.9%. OPT-350M in
229
+ fp16 under ONNX Runtime is bit-exact.
230
+
231
+ **Fold + keep, and it pays.** On a two-layer block whose attention is a custom
232
+ squared-ReLU flash kernel, KernelLens runs it in **7.63 ms**, versus **15.15 ms**
233
+ for the folded-only `torch.compile` formulation and **28.04 ms** for eager
234
+ folding — **1.99×** and **3.68×** faster. Against the unfused kernel path it is
235
+ **1.08×** faster than eager Triton (8.24 ms) and level with `torch.compile`
236
+ (7.55 ms). The engine folds the standard layers; the kernel keeps the operations
237
+ folding cannot express.
238
+
239
+ **Exact parity with native Triton** (`MaxDiff = 0.00e+00`) for LLaMA 3 RMSNorm,
240
+ PaLM SwiGLU, Qwen 2.5 RoPE, Liger fused cross-entropy, in-place mutation and the
241
+ quantized GEMMs, on both backends — and across all **fourteen** module patterns
242
+ of `tests/test_conformance.py`. Two multi-operator fusions are near-exact rather
243
+ than exact (`≤ 9.5e-07` residual+RMSNorm, `4.8e-07` Mamba's scan) and are
244
+ reported as such.
245
+
246
+ **Lowest VRAM** in a 4-layer LLaMA decoder: **464 MB** peak vs 478 MB (eager)
247
+ and 473 MB (`torch.compile`), with the fastest time-to-first-token (4.14 ms vs
248
+ 7.85/7.92 ms).
191
249
 
192
250
  ## Documentation
193
251
 
@@ -195,9 +253,11 @@ the [results page](https://just1truc.github.io/kernel_lens/#results).
195
253
  | :--- | :--- |
196
254
  | [Problem](https://just1truc.github.io/kernel_lens/#problem) | The five gaps a naive export hits |
197
255
  | [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, latency, end-to-end |
256
+ | [Results](https://just1truc.github.io/kernel_lens/#results) | Folding + kernels, parity, conformance, single-operator and end-to-end latency, **real pretrained models** |
257
+ | [Why an engine](https://just1truc.github.io/kernel_lens/#why) | `torch.compile` vs TensorRT, both directions |
199
258
  | [Install & API](https://just1truc.github.io/kernel_lens/#install) | Installation and API reference |
200
259
  | [Status](https://just1truc.github.io/kernel_lens/#status) | Roadmap |
260
+ | [FAQ](https://just1truc.github.io/kernel_lens/#faq) | Fair objections |
201
261
  | [Paper](https://just1truc.github.io/kernel_lens/#paper) | Abstract and citation |
202
262
 
203
263
  ## Reproduce the results
@@ -212,14 +272,30 @@ python tests/test_research_kernels.py # numerical parity
212
272
  python tests/benchmark_research_kernels.py # single-operator latency
213
273
  python tests/benchmark_e2e_llama.py # end-to-end decoder
214
274
  python tests/test_fp8_gemm.py # FP8 Tensor Core path
275
+ python tests/measure_codegen.py # generated C++ volume, per file
276
+ python tests/benchmark_real_model.py --model facebook/opt-1.3b # real OPT, both precisions
215
277
  ```
216
278
 
279
+ The real-model matrix (OPT-125M → 6.7B, fp32 and fp16) runs unattended on a fresh
280
+ CUDA host: `scripts/remote_setup.sh` provisions nvcc, TensorRT, `onnxruntime-gpu`
281
+ and `transformers` and installs the package; `scripts/remote_bench_all.sh` then
282
+ sweeps the matrix into `results/` and writes `summary.json`, recording failures
283
+ instead of aborting. See
284
+ [`docs/assets/data/a100/runpod_runbook.md`](docs/assets/data/a100/runpod_runbook.md) for the
285
+ exact invocation and the environment traps (CUDA 12.8's `nvcc` does not build
286
+ against a modern libstdc++; CUDA 13.3's does).
287
+
217
288
  Raw measurement JSON is committed under [`docs/assets/data/`](docs/assets/data/).
218
289
 
219
290
  ## Technical report
220
291
 
221
- The LaTeX source is [`architecture_report.tex`](architecture_report.tex); the TMLR
222
- variant is generated from it with `python3 scripts/make_tmlr.py`.
292
+ The paper source is [`kernel_lens_fresh.tex`](kernel_lens_fresh.tex). Two derived
293
+ variants are generated from it and must not be edited by hand:
294
+ `scripts/make_mlsys.py` produces the 2-column MLSys form
295
+ (`kernel_lens_mlsys.tex`) and `scripts/make_tmlr_fresh.py` produces the TMLR
296
+ forms — `kernel_lens_tmlr_fresh.tex` (named preprint) and
297
+ `kernel_lens_tmlr_submission.tex` (anonymized; **do not publish it alongside the
298
+ named preprint**).
223
299
 
224
300
  ```bibtex
225
301
  @article{duc2026kernellens,
@@ -227,7 +303,7 @@ variant is generated from it with `python3 scripts/make_tmlr.py`.
227
303
  author = {Duc, Justin},
228
304
  institution = {Galiad Research},
229
305
  year = {2026},
230
- version = {1.1.9},
306
+ version = {1.2.1},
231
307
  doi = {10.5281/zenodo.23194523},
232
308
  url = {https://doi.org/10.5281/zenodo.23194523}
233
309
  }
@@ -7,6 +7,9 @@ onnxruntime-gpu>=1.17.0
7
7
  tensorrt<11.0.0,>=8.6.0
8
8
  tensorrt-cu12-libs<11.0.0
9
9
 
10
+ [bench]
11
+ transformers>=4.40.0
12
+
10
13
  [ort]
11
14
  onnxruntime-gpu>=1.17.0
12
15
 
@@ -4,18 +4,36 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "kernel-lens"
7
- version = "1.2.0"
8
- description = "A production-grade, multi-backend compiler for Triton kernels (ONNX & TensorRT)"
7
+ version = "1.2.1"
8
+ description = "Compile PyTorch models with custom @triton.jit kernels into ONNX Runtime and TensorRT C++ plugins"
9
9
  readme = "README.md"
10
10
  license = { file = "LICENSE" }
11
11
  requires-python = ">=3.10"
12
12
  authors = [
13
- { name = "Compiler Engineer" }
13
+ { name = "Justin Duc", email = "justin.duc@galiad-research.com" }
14
+ ]
15
+ keywords = [
16
+ "triton",
17
+ "tensorrt",
18
+ "onnxruntime",
19
+ "compiler",
20
+ "gpu",
21
+ "cuda",
22
+ "kernel",
23
+ "inference",
24
+ "deep-learning",
14
25
  ]
15
26
  classifiers = [
27
+ "Development Status :: 4 - Beta",
28
+ "Intended Audience :: Science/Research",
16
29
  "Programming Language :: Python :: 3",
30
+ "Programming Language :: Python :: 3.10",
31
+ "Programming Language :: Python :: 3.11",
32
+ "Programming Language :: Python :: 3.12",
33
+ "Programming Language :: Python :: 3.13",
17
34
  "License :: OSI Approved :: MIT License",
18
35
  "Operating System :: POSIX :: Linux",
36
+ "Environment :: GPU :: NVIDIA CUDA",
19
37
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
20
38
  ]
21
39
 
@@ -31,10 +49,16 @@ dependencies = [
31
49
  ort = ["onnxruntime-gpu>=1.17.0"]
32
50
  trt = ["tensorrt>=8.6.0,<11.0.0", "tensorrt-cu12-libs<11.0.0"]
33
51
  all = ["onnxruntime-gpu>=1.17.0", "tensorrt>=8.6.0,<11.0.0", "tensorrt-cu12-libs<11.0.0"]
52
+ # Needed only by tests/benchmark_real_model.py, which serves real checkpoints.
53
+ bench = ["transformers>=4.40.0"]
34
54
 
35
55
  [project.urls]
56
+ Homepage = "https://just1truc.github.io/kernel_lens/"
57
+ Documentation = "https://just1truc.github.io/kernel_lens/"
58
+ Changelog = "https://github.com/Just1truc/kernel_lens/blob/main/CHANGELOG.md"
36
59
  "Source Code" = "https://github.com/Just1truc/kernel_lens"
37
60
  "Bug Tracker" = "https://github.com/Just1truc/kernel_lens/issues"
61
+ Paper = "https://doi.org/10.5281/zenodo.23194523"
38
62
 
39
63
  [tool.setuptools.packages.find]
40
64
  include = ["kernel_lens*"]
File without changes
File without changes