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.
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/PKG-INFO +105 -29
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/README.md +88 -26
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/__init__.py +1 -1
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/builder.py +17 -3
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/runtime/engine.py +95 -8
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/PKG-INFO +105 -29
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/requires.txt +3 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/pyproject.toml +27 -3
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/LICENSE +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/ort_gen.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/backends/trt_gen.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/ast_analyzer.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/core.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/interaction.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/manifest.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/onnx_exporter.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/compiler/tracer.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/config.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInfer.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferImpl.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferLegacyDims.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferPlugin.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferPluginBase.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferPluginUtils.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntime.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntimeBase.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntimeCommon.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferRuntimePlugin.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/include/tensorrt/NvInferVersion.h +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/runtime/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/utils/__init__.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/utils/deployment.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens/utils/env_check.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/SOURCES.txt +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/dependency_links.txt +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/kernel_lens.egg-info/top_level.txt +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/setup.cfg +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_alignment_guard.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_architectural_fixes.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_conformance.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_edge_cases.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_fp8_gemm.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_fused_cross_entropy.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_fused_residual_rms.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_inplace.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_inplace_dtypes.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_int4_dequant_gemm.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_mamba_selective_scan.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_research_kernels.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_standalone.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_standalone_launch.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_standalone_ptx.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_tricky_triton_kernels.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_triton_vs_torch_compile.py +0 -0
- {kernel_lens-1.2.0 → kernel_lens-1.2.1}/tests/test_trt.py +0 -0
- {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.
|
|
4
|
-
Summary:
|
|
5
|
-
Author:
|
|
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
|
|
117
|
-
>
|
|
118
|
-
>
|
|
119
|
-
>
|
|
120
|
-
>
|
|
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
|
-
|
|
176
|
-
|
|
177
|
-
the
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
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,
|
|
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
|
|
222
|
-
|
|
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
|
|
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
|
|
67
|
-
>
|
|
68
|
-
>
|
|
69
|
-
>
|
|
70
|
-
>
|
|
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
|
-
|
|
126
|
-
|
|
127
|
-
the
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
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,
|
|
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
|
|
172
|
-
|
|
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
|
|
242
|
+
version = {1.2.1},
|
|
181
243
|
doi = {10.5281/zenodo.23194523},
|
|
182
244
|
url = {https://doi.org/10.5281/zenodo.23194523}
|
|
183
245
|
}
|
|
@@ -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 + [
|
|
@@ -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
|
-
|
|
311
|
-
|
|
312
|
-
|
|
313
|
-
|
|
314
|
-
|
|
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
|
|
362
|
-
self._trt_engine = runtime
|
|
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.
|
|
4
|
-
Summary:
|
|
5
|
-
Author:
|
|
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
|
|
117
|
-
>
|
|
118
|
-
>
|
|
119
|
-
>
|
|
120
|
-
>
|
|
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
|
-
|
|
176
|
-
|
|
177
|
-
the
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
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,
|
|
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
|
|
222
|
-
|
|
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
|
|
306
|
+
version = {1.2.1},
|
|
231
307
|
doi = {10.5281/zenodo.23194523},
|
|
232
308
|
url = {https://doi.org/10.5281/zenodo.23194523}
|
|
233
309
|
}
|
|
@@ -4,18 +4,36 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "kernel-lens"
|
|
7
|
-
version = "1.2.
|
|
8
|
-
description = "
|
|
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 = "
|
|
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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|