kernel-lens 1.2.0__tar.gz → 1.3.0__tar.gz

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