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

{kernel_lens-1.3.0 → kernel_lens-2.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel-lens
-Version: 1.3.0
+Version: 2.0.1
 Summary: Compile PyTorch models with custom @triton.jit kernels into ONNX Runtime and TensorRT C++ plugins
 Author-email: Justin Duc <justin.duc@galiad-research.com>
 License: MIT License
@@ -52,12 +52,12 @@ Requires-Dist: onnx>=1.15.0
 Provides-Extra: ort
 Requires-Dist: onnxruntime-gpu>=1.17.0; extra == "ort"
 Provides-Extra: trt
-Requires-Dist: tensorrt<11.0.0,>=8.6.0; extra == "trt"
-Requires-Dist: tensorrt-cu12-libs<11.0.0; extra == "trt"
+Requires-Dist: tensorrt<12.0.0,>=8.6.0; extra == "trt"
+Requires-Dist: tensorrt-cu12-libs; extra == "trt"
 Provides-Extra: all
 Requires-Dist: onnxruntime-gpu>=1.17.0; extra == "all"
-Requires-Dist: tensorrt<11.0.0,>=8.6.0; extra == "all"
-Requires-Dist: tensorrt-cu12-libs<11.0.0; extra == "all"
+Requires-Dist: tensorrt<12.0.0,>=8.6.0; extra == "all"
+Requires-Dist: tensorrt-cu12-libs; extra == "all"
 Provides-Extra: bench
 Requires-Dist: transformers>=4.40.0; extra == "bench"
 Dynamic: license-file
@@ -86,7 +86,8 @@ Dynamic: license-file
 
 **KernelLens** compiles a PyTorch model that contains custom `@triton.jit` kernels
 into native C++ artifacts for **ONNX Runtime** (`OrtCustomOp`) and **NVIDIA
-TensorRT 10.x** (`nvinfer1::IPluginV2DynamicExt`). You write Triton in Python; it
+TensorRT 8.6 to 11** (`nvinfer1::IPluginV2DynamicExt`, and `nvinfer1::IPluginV3`
+on TensorRT 10 and later). You write Triton in Python; it
 emits the C++ plugin, builds the `.so`, and runs it with **zero-copy VRAM
 binding** — no manual C++ and no Python in the serving path.
 
@@ -114,6 +115,12 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
 - **Dynamic grids in C++** — launch-grid expressions are transpiled to C++
   integer math over runtime shapes, so one plugin serves dynamic batch and
   sequence lengths.
+- **One package, several TensorRT versions** — the plugin interface is chosen
+  from the environment. TensorRT 10 deprecated the whole `IPluginV2*` hierarchy
+  and TensorRT 11 removed it, so `plugin_api="auto"` (the default) emits
+  `IPluginV3` plugins where the runtime supports them and `IPluginV2DynamicExt`
+  plugins otherwise. This spans TensorRT 8.6 through 10.x today; TensorRT 11+
+  works when you point `TENSORRT_INCLUDE_DIR` at matching headers.
 - **Shape-generic engines** — `kl.compile(..., dynamic_shapes=True)` exports the
   leading dimension as a symbolic axis, so a single engine serves any size in a
   declared range via a TensorRT optimization profile. Verified on Blackwell
@@ -139,9 +146,12 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
 > conformance patterns and the quantized GEMMs, on both backends; two
 > multi-operator fusions land just above zero and are reported as such. Keeping
 > the kernel pays when the baseline leaves the graph unfused and breaks even when
-> it already fuses. The TensorRT-LLM backend is design-only (see
-> [`docs/assets/data/a100/trtllm_backend_design.md`](docs/assets/data/a100/trtllm_backend_design.md)),
-> and the quantized GEMMs are validated on ONNX Runtime only. Full methodology
+> it already fuses. A TensorRT-LLM backend is out of scope; the
+> [design note](docs/assets/data/a100/trtllm_backend_design.md) records what
+> integrating one would take. The kernels themselves are usable there — they are
+> ordinary `@triton.jit` functions, and an A100 measurement confirmed a KernelLens
+> RMSNorm is bit-identical to the flashinfer kernel TensorRT-LLM already uses
+> (`0.00e+00`). The quantized GEMMs are validated on ONNX Runtime only. Full methodology
 > and every raw JSON are under [`docs/assets/data/`](docs/assets/data/).
 
 ## Installation
@@ -157,11 +167,34 @@ pip install kernel-lens[all]     # both
 Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
 generator emits `IPluginV2DynamicExt`).
 
+**TensorRT versions.** Handled automatically. The pip TensorRT wheels ship **no
+headers**, so KernelLens fetches the set matching the installed runtime from the
+TensorRT OSS repository and caches it under `~/.kernel_lens_headers/<major>.<minor>`;
+when the headers it already has match the runtime major, nothing is fetched. If
+the download is unavailable the fallback is the bundled set, and a
+header/runtime major mismatch is then reported instead of crashing inside
+TensorRT. `plugin_api="auto"` picks the interface from the runtime:
+
+| TensorRT | `IPluginV2*` | `IPluginV3` | emitted |
+| :--- | :--- | :--- | :--- |
+| < 10 (8.6, 9.x) | present | absent | `v2` |
+| 10.x | deprecated | present | `v3` |
+| 11+ | **removed** | present | `v3` |
+
+Verified on TensorRT 10.16 and 11.3, both 16/16 on the conformance suite. Force
+a choice with `plugin_api="v2"` / `"v3"`, or the `KERNEL_LENS_PLUGIN_API`
+environment variable. To supply headers yourself — offline installs, or a TensorRT
+build the OSS repository does not track — set `TENSORRT_INCLUDE_DIR`, or
+`KERNEL_LENS_FETCH_HEADERS=0` to disable fetching entirely.
+
 **Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
 `.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
 compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
 specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
-`TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
+`TENSORRT_INCLUDE_DIR` (override the TensorRT headers),
+`KERNEL_LENS_HEADER_CACHE` (where fetched TensorRT headers are kept),
+`KERNEL_LENS_FETCH_HEADERS=0` (disable fetching),
+`KERNEL_LENS_PLUGIN_API` (`auto`, `v2` or `v3`), and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
 above 2 GB from disk rather than from a memory buffer (opt-in; see the
 [changelog](CHANGELOG.md)).
 

{kernel_lens-1.3.0 → kernel_lens-2.0.1}/README.md

@@ -22,7 +22,8 @@
 
 **KernelLens** compiles a PyTorch model that contains custom `@triton.jit` kernels
 into native C++ artifacts for **ONNX Runtime** (`OrtCustomOp`) and **NVIDIA
-TensorRT 10.x** (`nvinfer1::IPluginV2DynamicExt`). You write Triton in Python; it
+TensorRT 8.6 to 11** (`nvinfer1::IPluginV2DynamicExt`, and `nvinfer1::IPluginV3`
+on TensorRT 10 and later). You write Triton in Python; it
 emits the C++ plugin, builds the `.so`, and runs it with **zero-copy VRAM
 binding** — no manual C++ and no Python in the serving path.
 
@@ -50,6 +51,12 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
 - **Dynamic grids in C++** — launch-grid expressions are transpiled to C++
   integer math over runtime shapes, so one plugin serves dynamic batch and
   sequence lengths.
+- **One package, several TensorRT versions** — the plugin interface is chosen
+  from the environment. TensorRT 10 deprecated the whole `IPluginV2*` hierarchy
+  and TensorRT 11 removed it, so `plugin_api="auto"` (the default) emits
+  `IPluginV3` plugins where the runtime supports them and `IPluginV2DynamicExt`
+  plugins otherwise. This spans TensorRT 8.6 through 10.x today; TensorRT 11+
+  works when you point `TENSORRT_INCLUDE_DIR` at matching headers.
 - **Shape-generic engines** — `kl.compile(..., dynamic_shapes=True)` exports the
   leading dimension as a symbolic axis, so a single engine serves any size in a
   declared range via a TensorRT optimization profile. Verified on Blackwell
@@ -75,9 +82,12 @@ compiled_model.run(inputs, backend="onnx"|"tensorrt")
 > conformance patterns and the quantized GEMMs, on both backends; two
 > multi-operator fusions land just above zero and are reported as such. Keeping
 > the kernel pays when the baseline leaves the graph unfused and breaks even when
-> it already fuses. The TensorRT-LLM backend is design-only (see
-> [`docs/assets/data/a100/trtllm_backend_design.md`](docs/assets/data/a100/trtllm_backend_design.md)),
-> and the quantized GEMMs are validated on ONNX Runtime only. Full methodology
+> it already fuses. A TensorRT-LLM backend is out of scope; the
+> [design note](docs/assets/data/a100/trtllm_backend_design.md) records what
+> integrating one would take. The kernels themselves are usable there — they are
+> ordinary `@triton.jit` functions, and an A100 measurement confirmed a KernelLens
+> RMSNorm is bit-identical to the flashinfer kernel TensorRT-LLM already uses
+> (`0.00e+00`). The quantized GEMMs are validated on ONNX Runtime only. Full methodology
 > and every raw JSON are under [`docs/assets/data/`](docs/assets/data/).
 
 ## Installation
@@ -93,11 +103,34 @@ pip install kernel-lens[all]     # both
 Triton ≥ 2.1, and `nvcc` + `g++` on `PATH`. TensorRT must be `< 11` (the
 generator emits `IPluginV2DynamicExt`).
 
+**TensorRT versions.** Handled automatically. The pip TensorRT wheels ship **no
+headers**, so KernelLens fetches the set matching the installed runtime from the
+TensorRT OSS repository and caches it under `~/.kernel_lens_headers/<major>.<minor>`;
+when the headers it already has match the runtime major, nothing is fetched. If
+the download is unavailable the fallback is the bundled set, and a
+header/runtime major mismatch is then reported instead of crashing inside
+TensorRT. `plugin_api="auto"` picks the interface from the runtime:
+
+| TensorRT | `IPluginV2*` | `IPluginV3` | emitted |
+| :--- | :--- | :--- | :--- |
+| < 10 (8.6, 9.x) | present | absent | `v2` |
+| 10.x | deprecated | present | `v3` |
+| 11+ | **removed** | present | `v3` |
+
+Verified on TensorRT 10.16 and 11.3, both 16/16 on the conformance suite. Force
+a choice with `plugin_api="v2"` / `"v3"`, or the `KERNEL_LENS_PLUGIN_API`
+environment variable. To supply headers yourself — offline installs, or a TensorRT
+build the OSS repository does not track — set `TENSORRT_INCLUDE_DIR`, or
+`KERNEL_LENS_FETCH_HEADERS=0` to disable fetching entirely.
+
 **Environment:** `KERNEL_LENS_CACHE_DIR` (put it on a persistent volume so the
 `.so` / `.engine` artifacts survive `kl.load`), `KERNEL_LENS_CCBIN` (host
 compiler for `nvcc` — CUDA toolkits before 12.9 reject GCC 13's exception
 specifiers, so set it to `g++-12` there), `KERNEL_LENS_DEBUG` (or `KL_VERBOSE`),
-`TENSORRT_INCLUDE_DIR`, and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
+`TENSORRT_INCLUDE_DIR` (override the TensorRT headers),
+`KERNEL_LENS_HEADER_CACHE` (where fetched TensorRT headers are kept),
+`KERNEL_LENS_FETCH_HEADERS=0` (disable fetching),
+`KERNEL_LENS_PLUGIN_API` (`auto`, `v2` or `v3`), and `KERNEL_LENS_STREAM_READER=1` to deserialize engines
 above 2 GB from disk rather than from a memory buffer (opt-in; see the
 [changelog](CHANGELOG.md)).
 

{kernel_lens-1.3.0 → kernel_lens-2.0.1}/kernel_lens/__init__.py

@@ -7,7 +7,7 @@ from .config import set_verbose, is_verbose
 
 
 
-__version__ = "1.3.0"
+__version__ = "2.0.0"
 
 __all__ = [
     "compile",

{kernel_lens-1.3.0 → kernel_lens-2.0.1}/kernel_lens/backends/builder.py

@@ -249,6 +249,21 @@ def build_trt_plugin(trt_plugins_dir: str, cache_dir: str):
             arch_flag = "-gencode=arch=compute_75,code=sm_75"
 
         trt_inc_dirs = []
+
+        # Headers matching the installed TensorRT, fetched and cached on first
+        # use when the bundled set is a different major version. See
+        # utils/trt_headers.py for why this is necessary: the pip TensorRT
+        # wheels ship no headers at all.
+        try:
+            from ..utils.trt_headers import resolve as _resolve_headers
+            _hdr_dir, _hdr_reason = _resolve_headers()
+            from ..config import debug_print as _dp
+            _dp(f"[KernelLens] TensorRT headers: {_hdr_dir or 'bundled'} ({_hdr_reason})")
+            if _hdr_dir:
+                trt_inc_dirs.append(_hdr_dir)
+        except Exception:
+            pass
+
         if os.environ.get("TENSORRT_INCLUDE_DIR"):
             trt_inc_dirs.append(os.environ["TENSORRT_INCLUDE_DIR"])
         user_trt_inc = os.path.expanduser("~/tensorrt_headers")

{kernel_lens-1.3.0 → kernel_lens-2.0.1}/kernel_lens/backends/trt_gen.py

@@ -28,37 +28,31 @@ class TensorRTPluginGenerator:
         self.plugin_namespace = plugin_namespace
         self.plugin_version = plugin_version
         
-    def _generate_kernel_h(self, manifest: KernelManifest) -> str:
-        plugin_name = f"{manifest.kernel_name}Plugin"
-
-        constant_decls = []
-        for arg in manifest.arguments:
-            if arg.kind == 'scalar':
-                ctype = "float" if isinstance(arg.value, float) else "int"
-                constant_decls.append(f"{ctype} m_{arg.name};")
-        dynamic_members_cpp = "\n    ".join(constant_decls)
-
-        def get_trt_datatype(dtype_obj):
-            dtype_str = str(dtype_obj).lower()
-            if 'float8' in dtype_str or 'fp8' in dtype_str or 'e4m3' in dtype_str or 'e5m2' in dtype_str:
-                return "nvinfer1::DataType::kFP8"
-            elif 'bool' in dtype_str:
-                return "nvinfer1::DataType::kBOOL"
-            elif 'uint8' in dtype_str or 'int8' in dtype_str or 'int4' in dtype_str or 'uint4' in dtype_str:
-                return "nvinfer1::DataType::kINT8"
-            elif 'bfloat16' in dtype_str or 'bf16' in dtype_str:
-                return "nvinfer1::DataType::kBF16"
-            elif 'float16' in dtype_str or 'half' in dtype_str:
-                return "nvinfer1::DataType::kHALF"
-            elif 'float' in dtype_str:
-                return "nvinfer1::DataType::kFLOAT"
-            elif 'int64' in dtype_str or 'long' in dtype_str:
-                return "nvinfer1::DataType::kINT64"
-            elif 'int' in dtype_str:
-                return "nvinfer1::DataType::kINT32"
-            else:
-                return "nvinfer1::DataType::kFLOAT"
-
+    @staticmethod
+    def trt_datatype(dtype_obj: str) -> str:
+        """Map a torch dtype onto the TensorRT DataType."""
+        dtype_str = str(dtype_obj).lower()
+        if 'float8' in dtype_str or 'fp8' in dtype_str or 'e4m3' in dtype_str or 'e5m2' in dtype_str:
+            return "nvinfer1::DataType::kFP8"
+        elif 'bool' in dtype_str:
+            return "nvinfer1::DataType::kBOOL"
+        elif 'uint8' in dtype_str or 'int8' in dtype_str or 'int4' in dtype_str or 'uint4' in dtype_str:
+            return "nvinfer1::DataType::kINT8"
+        elif 'bfloat16' in dtype_str or 'bf16' in dtype_str:
+            return "nvinfer1::DataType::kBF16"
+        elif 'float16' in dtype_str or 'half' in dtype_str:
+            return "nvinfer1::DataType::kHALF"
+        elif 'float' in dtype_str:
+            return "nvinfer1::DataType::kFLOAT"
+        elif 'int64' in dtype_str or 'long' in dtype_str:
+            return "nvinfer1::DataType::kINT64"
+        elif 'int' in dtype_str:
+            return "nvinfer1::DataType::kINT32"
+        else:
+            return "nvinfer1::DataType::kFLOAT"
+    @staticmethod
+    def supports_format_body(manifest: KernelManifest, tensor_prefix: str = "") -> str:
+        """The supportsFormatCombination body, shared by both plugin emitters."""
         in_args = [a for a in manifest.arguments if a.kind in ('input', 'inplace')]
 
         # Enforce the dtype the kernel was compiled for, position by position.
@@ -67,7 +61,7 @@ class TensorRTPluginGenerator:
         _ordered_in = [a for a in manifest.arguments if a.kind in ('input', 'inplace')]
         _ordered_sc = [a for a in manifest.arguments if a.kind == 'scalar' and not getattr(a, 'is_constexpr', False)]
         _ordered_out = [a for a in manifest.arguments if a.kind in ('output', 'inplace')]
-        _pos_types = [get_trt_datatype(a.dtype) for a in (_ordered_in + _ordered_sc + _ordered_out)]
+        _pos_types = [TensorRTPluginGenerator.trt_datatype(a.dtype) for a in (_ordered_in + _ordered_sc + _ordered_out)]
         _pos_is_tensor = ([True] * len(_ordered_in) + [False] * len(_ordered_sc) + [True] * len(_ordered_out))
         _types_cpp = ", ".join(_pos_types) if _pos_types else "nvinfer1::DataType::kFLOAT"
         _is_tensor_cpp = ", ".join("true" if t else "false" for t in _pos_is_tensor) if _pos_is_tensor else "true"
@@ -96,7 +90,24 @@ class TensorRTPluginGenerator:
             )
         else:
             supports_format_cxx = type_guard_cxx + "return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;"
-        
+        # IPluginV3 hands supportsFormatCombination a DynamicPluginTensorDesc,
+        # whose fields live under .desc; V2 hands it a PluginTensorDesc directly.
+        if tensor_prefix:
+            supports_format_cxx = supports_format_cxx.replace("inOut[pos].", f"inOut[pos]{tensor_prefix}.")
+        return supports_format_cxx
+
+    def _generate_kernel_h(self, manifest: KernelManifest) -> str:
+        plugin_name = f"{manifest.kernel_name}Plugin"
+
+        constant_decls = []
+        for arg in manifest.arguments:
+            if arg.kind == 'scalar':
+                ctype = "float" if isinstance(arg.value, float) else "int"
+                constant_decls.append(f"{ctype} m_{arg.name};")
+        dynamic_members_cpp = "\n    ".join(constant_decls)
+
+        get_trt_datatype = TensorRTPluginGenerator.trt_datatype
+        supports_format_cxx = TensorRTPluginGenerator.supports_format_body(manifest)
         output_dim_cases = []
         out_args = [a for a in manifest.arguments if a.kind in ('output', 'inplace')]
         _dyn_axes = tuple(getattr(manifest, "dynamic_axes", ()) or ())
@@ -217,12 +228,15 @@ private:
 '''
         return textwrap.dedent(tpl).strip()
 
-    def _generate_kernel_cu(self, manifest: KernelManifest) -> str:
-        plugin_name = f"{manifest.kernel_name}Plugin"
-        ptx_str = manifest.ptx
-        ptx_encoded = json.dumps(ptx_str)
-        block_size = (manifest.num_warps if manifest.num_warps > 0 else 4) * 32
-        
+    @staticmethod
+    def build_shared_body(manifest: KernelManifest, block_size: int):
+        """Build the pieces the enqueue implementation needs.
+
+        Shared verbatim between the IPluginV2DynamicExt and the
+        IPluginV3 emitters: the PTX entry name, the kernelParams
+        assembly and the output-type dispatch are the same; only the
+        class hierarchy and the method signatures around them differ.
+        """
         scalars = [arg for arg in manifest.arguments if arg.kind == 'scalar']
         nb_outputs = sum(1 for arg in manifest.arguments if arg.kind == 'output')
         nb_outputs = max(1, nb_outputs)
@@ -388,10 +402,41 @@ private:
                 
             output_type_lines.append(f"if (index == {i}) return {trt_type};")
             
-        dynamic_output_types = "\n    ".join(output_type_lines)
+        return {
+            "ptx_entry_name": ptx_entry_name,
+            "dynamic_args_cpp": "\n    ".join(arg_setup_lines),
+            "dynamic_output_types": "\n    ".join(output_type_lines),
+            "grid_x_cxx": grid_x_cxx,
+            "grid_y_cxx": grid_y_cxx,
+            "grid_z_cxx": grid_z_cxx,
+            "num_ptx_slots": num_ptx_slots,
+            "init_str": init_str,
+            "size_additions": size_additions,
+            "deserialize_cpp": deserialize_cpp,
+            "serialize_cpp": serialize_cpp,
+            "clone_cpp": clone_cpp,
+        }
 
-        dynamic_args_cpp = "\n    ".join(arg_setup_lines)
+    def _generate_kernel_cu(self, manifest: KernelManifest) -> str:
+        plugin_name = f"{manifest.kernel_name}Plugin"
+        ptx_str = manifest.ptx
+        ptx_encoded = json.dumps(ptx_str)
+        block_size = (manifest.num_warps if manifest.num_warps > 0 else 4) * 32
         
+        shared = TensorRTPluginGenerator.build_shared_body(manifest, block_size)
+        ptx_entry_name = shared["ptx_entry_name"]
+        dynamic_args_cpp = shared["dynamic_args_cpp"]
+        dynamic_output_types = shared["dynamic_output_types"]
+        grid_x_cxx = shared["grid_x_cxx"]
+        grid_y_cxx = shared["grid_y_cxx"]
+        grid_z_cxx = shared["grid_z_cxx"]
+        num_ptx_slots = shared["num_ptx_slots"]
+        init_str = shared["init_str"]
+        size_additions = shared["size_additions"]
+        deserialize_cpp = shared["deserialize_cpp"]
+        serialize_cpp = shared["serialize_cpp"]
+        clone_cpp = shared["clone_cpp"]
+
         tpl = f'''
 #include "{manifest.kernel_name}Plugin.h"
 #include <cuda.h>