chore: import upstream snapshot with attribution

tokenspeed-kernel/python/tokenspeed_kernel/plugins/README.md

@@ -0,0 +1,189 @@
+# Tokenspeed-kernel Plugin System
+
+> **Status: experimental.** The plugin contract — entry-point group name,
+> `register()` signature, `KernelSpec` fields, selection priority semantics,
+> and the `tokenspeed_kernel.plugins` Python API — may change without
+> backwards-compatibility guarantees while we shake the design out. Pin to
+> an exact `tokenspeed-kernel` version in your plugin's dependencies.
+
+This subpackage lets third-party packages register kernel implementations
+into the global `KernelRegistry` without modifying `tokenspeed-kernel`
+itself.
+
+## How it works
+
+1. Your plugin package exposes a `register()` function that calls
+   `tokenspeed_kernel.registry.register_kernel(...)` (or
+   `KernelRegistry.get().register(...)`) for each kernel it provides.
+2. Your `pyproject.toml` advertises that function under the
+   `tokenspeed_kernel.plugins` entry-point group.
+3. The host application (engine, benchmark, notebook, etc.) calls
+   `tokenspeed_kernel.plugins.discover_plugins()` once at startup, after
+   built-in kernels have been imported. Discovery walks the entry-point
+   group and invokes each `register()`.
+
+Loading is **fully explicit** — importing `tokenspeed_kernel` or
+`tokenspeed_kernel.plugins` does **not** trigger discovery on its own.
+
+## Example plugin
+
+A minimal out-of-tree package that contributes a custom decode-attention
+kernel for NVIDIA Hopper.
+
+```
+my-kernels-plugin/
+├── pyproject.toml
+└── my_kernels_plugin/
+    └── __init__.py
+```
+
+`my_kernels_plugin/__init__.py`:
+
+```python
+import torch
+
+from tokenspeed_kernel.platform import ArchVersion, CapabilityRequirement
+from tokenspeed_kernel.signature import format_signatures
+from tokenspeed_kernel.registry import register_kernel
+
+
+def register() -> None:
+    """Entry point invoked by tokenspeed_kernel.plugins.discover_plugins()."""
+
+    @register_kernel(
+        "attention",
+        "decode",
+        solution="my_custom",
+        signatures=format_signatures(
+            ("q", "k_cache", "v_cache"), "dense", {torch.bfloat16}
+        ),
+        capability=CapabilityRequirement(
+            vendors=frozenset({"nvidia"}),
+            min_arch_version=ArchVersion(9, 0),
+        ),
+        # Built-in FlashInfer decode is priority 18; pick 19 to win selection.
+        priority=19,
+    )
+    def my_custom_attn_decode(q, kv_cache, page_table, seq_lens, **kwargs):
+        ...
+```
+
+`pyproject.toml`:
+
+```toml
+[project]
+name = "my-kernels-plugin"
+version = "0.1.0"
+# Pin tightly while the plugin contract is experimental.
+dependencies = ["tokenspeed-kernel==<exact-version>"]
+
+[project.entry-points."tokenspeed_kernel.plugins"]
+my_plugin = "my_kernels_plugin:register"
+```
+
+Install and load:
+
+```bash
+pip install -e ./my-kernels-plugin
+```
+
+```python
+import tokenspeed_kernel  # registers built-in kernels
+from tokenspeed_kernel.plugins import discover_plugins, list_plugins
+
+discover_plugins()
+print(list_plugins())  # -> [PluginInfo(name='my_plugin', ...)]
+```
+
+## Host-application integration
+
+Engines and other long-running hosts should call `discover_plugins()`
+exactly once at startup, after built-in kernel modules have been imported
+(so plugins can override built-ins by registering at a higher priority).
+
+```python
+import tokenspeed_kernel  # noqa: F401  -- registers built-ins
+from tokenspeed_kernel.plugins import discover_plugins
+
+discover_plugins()
+```
+
+For ad-hoc use (notebooks, scripts, tests), there is no need to use entry
+points at all — call `register_kernel(...)` directly:
+
+```python
+import torch
+from tokenspeed_kernel.signature import format_signatures
+from tokenspeed_kernel.registry import register_kernel
+
+
+@register_kernel(
+    "gemm",
+    "mm",
+    solution="experiment",
+    signatures=format_signatures(("a", "b"), "dense", {torch.bfloat16}),
+    priority=15,
+)
+def my_experimental_gemm(a, b, **kwargs):
+    ...
+```
+
+## Disabling plugins
+
+Plugins can be skipped without uninstalling them:
+
+```bash
+TOKENSPEED_KERNEL_DISABLE_PLUGINS="my_plugin,other_plugin" python ...
+```
+
+```python
+from tokenspeed_kernel.plugins import disable_plugin, discover_plugins
+
+disable_plugin("my_plugin")
+discover_plugins()
+```
+
+The names refer to entry-point names (the left-hand side of the
+`[project.entry-points."tokenspeed_kernel.plugins"]` table), not
+distribution names.
+
+## Inspection
+
+```bash
+python -m tokenspeed_kernel.plugins list
+python -m tokenspeed_kernel.plugins info my_plugin
+```
+
+```python
+from tokenspeed_kernel.plugins import list_plugins
+
+for info in list_plugins():
+    print(info.name, info.version, info.kernel_names)
+```
+
+## Selection contract
+
+- Priority is an integer in `[0, 20)`. Higher wins. The reference
+  implementation lives at `0`. Built-in optimized kernels typically sit at
+  `10`–`18`. Plugin authors who want to override a built-in should choose
+  a value strictly higher than the built-in they replace.
+- `discover_plugins()` walks entry points in alphabetical order by
+  entry-point name. When two registrations land at the same priority for
+  the same `(family, mode)`, the warning is emitted and selection becomes
+  load-order-dependent — set explicit, distinct priorities to avoid this.
+- A plugin whose `register()` raises does not crash discovery; a
+  `UserWarning` is emitted and other plugins continue loading.
+
+## Failure modes worth knowing
+
+- **No discovery call → no plugin kernels.** Forgetting to call
+  `discover_plugins()` is silent.
+- **Plugin loaded before built-ins.** If you call `discover_plugins()`
+  before `import tokenspeed_kernel`, plugins that intend to override
+  built-ins will appear to win, but the built-in modules will be imported
+  later and may overwrite the plugin's slot. Always import
+  `tokenspeed_kernel` first.
+- **Stale registry.** Calling `KernelRegistry.reset()` clears registered
+  kernels but leaves `_loaded_plugins` populated; re-discovery will skip
+  already-loaded plugins. Use `discover_plugins(force=True)` after a
+  reset.

tokenspeed-kernel/python/tokenspeed_kernel/plugins/__init__.py

@@ -0,0 +1,208 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Out-of-tree kernel plugin system.
+
+Third-party packages register kernel implementations via Python entry points
+in the ``tokenspeed_kernel.plugins`` group. Loading is explicit: the host
+application must call :func:`discover_plugins` (typically once at startup,
+after built-in kernels have been imported) for installed plugins to take
+effect. See ``README.md`` for details.
+"""
+
+from __future__ import annotations
+
+import importlib.metadata as importlib_metadata
+import logging
+import os
+import warnings
+from dataclasses import dataclass, field
+
+from tokenspeed_kernel.registry import KernelRegistry
+
+logger = logging.getLogger(__name__)
+
+ENTRY_POINT_GROUP = "tokenspeed_kernel.plugins"
+DISABLE_ENV_VAR = "TOKENSPEED_KERNEL_DISABLE_PLUGINS"
+
+__all__ = [
+    "PluginInfo",
+    "ENTRY_POINT_GROUP",
+    "DISABLE_ENV_VAR",
+    "discover_plugins",
+    "list_plugins",
+    "disable_plugin",
+    "enable_plugin",
+    "reset_plugins",
+]
+
+
+@dataclass(frozen=True)
+class PluginInfo:
+    """Metadata describing a discovered plugin."""
+
+    name: str
+    package: str
+    version: str
+    num_kernels: int
+    kernel_names: tuple[str, ...] = field(default_factory=tuple)
+
+
+_loaded_plugins: dict[str, PluginInfo] = {}
+_disabled_plugins: set[str] = set()
+
+
+def _disabled_from_env() -> set[str]:
+    raw = os.environ.get(DISABLE_ENV_VAR, "")
+    return {part.strip() for part in raw.split(",") if part.strip()}
+
+
+def disable_plugin(name: str) -> None:
+    """Mark a plugin as disabled. Future ``discover_plugins`` calls skip it."""
+    _disabled_plugins.add(name)
+
+
+def enable_plugin(name: str) -> None:
+    """Remove a plugin from the disabled set (does not auto-load it)."""
+    _disabled_plugins.discard(name)
+
+
+def reset_plugins() -> None:
+    """Clear all plugin tracking state (for testing)."""
+    _loaded_plugins.clear()
+    _disabled_plugins.clear()
+
+
+def list_plugins() -> list[PluginInfo]:
+    """Return metadata for all plugins loaded in this process."""
+    return list(_loaded_plugins.values())
+
+
+def _entry_points(group: str):
+    """Compatibility shim across importlib.metadata API variants."""
+    eps = importlib_metadata.entry_points()
+    # Python 3.10+ exposes a SelectableGroups API.
+    selector = getattr(eps, "select", None)
+    if callable(selector):
+        return list(selector(group=group))
+    return list(eps.get(group, []))  # pragma: no cover - legacy path
+
+
+def _resolve_package_and_version(ep) -> tuple[str, str]:
+    """Best-effort distribution name and version for an entry point."""
+    dist = getattr(ep, "dist", None)
+    if dist is not None:
+        meta = getattr(dist, "metadata", None)
+        if meta is not None:
+            name = meta.get("Name") or ""
+        else:
+            name = getattr(dist, "name", "") or ""
+        version = getattr(dist, "version", "") or ""
+        if name:
+            return name, version
+    # Fallback: derive from the entry point's module path.
+    package = (ep.value or "").split(":", 1)[0].split(".", 1)[0]
+    version = ""
+    if package:
+        try:
+            version = importlib_metadata.version(package)
+        except importlib_metadata.PackageNotFoundError:
+            version = ""
+    return package, version
+
+
+def _new_kernel_names(before: set[str], after: set[str]) -> tuple[str, ...]:
+    return tuple(sorted(after - before))
+
+
+def _check_priority_collisions(
+    registry: KernelRegistry, new_names: tuple[str, ...]
+) -> None:
+    for name in new_names:
+        spec = registry.get_by_name(name)
+        if spec is None:
+            continue
+        peers = registry.get_for_operator(spec.family, spec.mode)
+        same = [s for s in peers if s.priority == spec.priority and s.name != name]
+        for other in same:
+            warnings.warn(
+                f"Plugin kernel '{name}' has equal priority ({spec.priority}) with "
+                f"existing kernel '{other.name}' for {spec.family}.{spec.mode}. "
+                "Set explicit priorities to make selection deterministic.",
+                stacklevel=3,
+            )
+
+
+def discover_plugins(*, force: bool = False) -> list[PluginInfo]:
+    """Discover and load installed kernel plugins via entry points.
+
+    Plugins disabled via :func:`disable_plugin` or the
+    ``TOKENSPEED_KERNEL_DISABLE_PLUGINS`` env var are skipped.
+
+    Already-loaded plugins are skipped unless ``force=True``, in which case
+    their ``register()`` is invoked again (useful for tests that reset the
+    registry).
+    """
+    disabled = _disabled_plugins | _disabled_from_env()
+    registry = KernelRegistry.get()
+    loaded: list[PluginInfo] = []
+
+    # Sort entry points alphabetically for deterministic load order. Plugins
+    # registering with the same priority will then have a stable winner.
+    eps = sorted(_entry_points(ENTRY_POINT_GROUP), key=lambda ep: ep.name)
+    for ep in eps:
+        if ep.name in disabled:
+            logger.info("Skipping disabled kernel plugin %r", ep.name)
+            continue
+        if ep.name in _loaded_plugins and not force:
+            continue
+
+        before = set(registry._by_name.keys())
+        try:
+            register_fn = ep.load()
+            register_fn()
+        except Exception as exc:
+            warnings.warn(
+                f"Failed to load kernel plugin {ep.name!r}: {exc}",
+                stacklevel=2,
+            )
+            continue
+
+        new_names = _new_kernel_names(before, set(registry._by_name.keys()))
+        _check_priority_collisions(registry, new_names)
+        package, version = _resolve_package_and_version(ep)
+        info = PluginInfo(
+            name=ep.name,
+            package=package,
+            version=version,
+            num_kernels=len(new_names),
+            kernel_names=new_names,
+        )
+        _loaded_plugins[ep.name] = info
+        loaded.append(info)
+        logger.info(
+            "Loaded kernel plugin %r (%s %s) registering %d kernels",
+            ep.name,
+            package,
+            version,
+            len(new_names),
+        )
+
+    return loaded

tokenspeed-kernel/python/tokenspeed_kernel/plugins/__main__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+from tokenspeed_kernel.plugins.cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tokenspeed-kernel/python/tokenspeed_kernel/plugins/cli.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from __future__ import annotations
+
+import argparse
+import sys
+from typing import Sequence
+
+from tokenspeed_kernel.plugins import discover_plugins, list_plugins
+from tokenspeed_kernel.registry import KernelRegistry
+
+
+def _cmd_list(_: argparse.Namespace) -> int:
+    plugins = sorted(list_plugins(), key=lambda p: p.name)
+    if not plugins:
+        print("No kernel plugins installed.")
+        return 0
+    print("Installed kernel plugins:")
+    for p in plugins:
+        version = f" ({p.version})" if p.version else ""
+        head = ", ".join(p.kernel_names[:4])
+        more = ", ..." if len(p.kernel_names) > 4 else ""
+        print(f"  {p.name}{version} — {p.num_kernels} kernels [{head}{more}]")
+    return 0
+
+
+def _cmd_info(args: argparse.Namespace) -> int:
+    plugins = {p.name: p for p in list_plugins()}
+    info = plugins.get(args.name)
+    if info is None:
+        print(f"Plugin {args.name!r} not loaded.", file=sys.stderr)
+        return 1
+
+    version = f" ({info.version})" if info.version else ""
+    print(f"Plugin: {info.name}{version}")
+    if info.package:
+        print(f"Package: {info.package}")
+    print("Kernels:")
+    registry = KernelRegistry.get()
+    for kname in info.kernel_names:
+        spec = registry.get_by_name(kname)
+        if spec is None:
+            print(f"  {kname:32s} (no longer registered)")
+            continue
+        op = f"{spec.family}.{spec.mode}"
+        cap = str(spec.capability) if spec.capability else "any"
+        print(f"  {kname:32s} {op:20s} priority={spec.priority}  {cap}")
+    return 0
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="python -m tokenspeed_kernel.plugins")
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    list_p = sub.add_parser("list", help="List installed kernel plugins")
+    list_p.set_defaults(func=_cmd_list)
+
+    info_p = sub.add_parser("info", help="Show details for a single plugin")
+    info_p.add_argument("name", help="Plugin entry point name")
+    info_p.set_defaults(func=_cmd_info)
+
+    # Plugin loading is explicit; discover here so the CLI sees them.
+    discover_plugins()
+
+    args = parser.parse_args(argv)
+    return args.func(args)