chore: import upstream snapshot with attribution

python/sglang/kernels/selector.py

@@ -0,0 +1,78 @@
+"""Fixed-path kernel resolution over the :data:`registry`.
+
+There is no priority ranking or heuristic backend selection. Each operator has
+a fixed call path — its :attr:`KernelSpec.target`:
+
+- an op with a single registered backend resolves to it directly;
+- an op with several registered backends must be resolved by naming the backend
+  explicitly (``backend=...``). The extra backends exist only as inventory, and
+  are never silently auto-picked.
+
+:func:`get_kernel` is the fast path used by the public ``ops.*`` wrappers: it
+resolves the spec to its callable and caches the result so repeated calls do
+not re-run resolution or re-import.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Callable, Optional
+
+from sglang.kernels.registry import registry
+from sglang.kernels.spec import KernelBackend, KernelSpec
+
+
+def select_kernel(op: str, backend: Optional[KernelBackend] = None) -> KernelSpec:
+    """Return the :class:`KernelSpec` for ``op`` (its fixed call path).
+
+    Parameters
+    ----------
+    op:
+        Operator id, ``"<group>.<name>"``.
+    backend:
+        Required only when ``op`` has more than one registered backend; selects
+        which one. For single-backend ops it is optional.
+
+    Raises
+    ------
+    KeyError
+        If ``op`` is unknown, or if ``backend`` is requested but not registered.
+    ValueError
+        If ``op`` has multiple backends and ``backend`` is not given.
+    """
+    specs = registry.get(op)
+    if not specs:
+        raise KeyError(f"No kernels registered for op {op!r}")
+
+    if backend is not None:
+        for spec in specs:
+            if spec.backend == backend:
+                return spec
+        raise KeyError(f"No '{backend.value}' backend registered for op {op!r}")
+
+    if len(specs) == 1:
+        return specs[0]
+
+    raise ValueError(
+        f"op {op!r} has multiple registered backends "
+        f"({[s.backend.value for s in specs]}); pass backend=... to choose one"
+    )
+
+
+@lru_cache(maxsize=None)
+def _resolve(op: str, backend: Optional[KernelBackend]) -> Callable:
+    return select_kernel(op, backend=backend).load()
+
+
+def get_kernel(op: str, backend: Optional[KernelBackend] = None) -> Callable:
+    """Resolve ``op`` to a callable kernel and cache it.
+
+    This is what the public ``sglang.kernels.ops.*`` wrappers call. The first
+    call resolves and imports the backend; later calls hit the cache.
+    """
+    return _resolve(op, backend)
+
+
+def clear_cache() -> None:
+    """Drop the resolved-callable cache (used by tests)."""
+    _resolve.cache_clear()