# Copyright 2025 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ Kernel registry for OpSlot-based dispatch. Provides a global registry of kernel implementations keyed by (op_name, variant, impl_name). Each kernel is described by a KernelSpec that includes a lazy factory, hardware requirements, and metadata. """ from __future__ import annotations from dataclasses import dataclass from typing import Callable from ..utils import logging from ..utils.device import IS_CUDA_AVAILABLE, IS_NPU_AVAILABLE, get_gpu_compute_capability logger = logging.get_logger(__name__) @dataclass(frozen=True) class HardwareRequirement: """Describes hardware constraints for a kernel.""" device_type: str # "gpu" | "npu" min_compute_capability: int | None = None # e.g. 70, 80, 90 # Inclusive upper bound for kernels that don't yet support newer arches # (e.g. FlashQLA today only ships SM90; SM100/SM120 wheels are WIP per # https://github.com/QwenLM/FlashQLA/issues/2). Drop this once the kernel # adds forward-compatibility for higher arches. max_compute_capability: int | None = None def is_satisfied(self) -> bool: if self.device_type == "gpu": if not IS_CUDA_AVAILABLE: return False cc = get_gpu_compute_capability() if self.min_compute_capability is not None and cc < self.min_compute_capability: return False if self.max_compute_capability is not None and cc > self.max_compute_capability: return False return True if self.device_type == "npu": # IS_NPU_AVAILABLE == is_torch_npu_available(): requires both the # torch_npu package AND an actual NPU device (unlike a bare import # check, which passes on dev boxes that merely have the library). return IS_NPU_AVAILABLE if self.device_type == "any": # Hardware-agnostic kernel (pure PyTorch). Used e.g. by chunk_loss # (F.linear + eager_cross_entropy in a chunked autograd Function), # which has no device-specific calls. Always satisfied — including # on CPU-only hosts (unit tests, weight materialization, dev boxes # without an accelerator). return True raise ValueError(f"Unknown device_type: {self.device_type!r} (expected 'gpu' | 'npu' | 'any')") @dataclass(frozen=True) class KernelSpec: """Describes a single kernel implementation registered under an op/variant. Attributes: name: Identifier exposed to users via the matching ``OpsImplementationConfig`` field (e.g. ``"liger_kernel"``, ``"triton"``, ``"quack"``). Must be unique within a given ``(op_name, variant)`` bucket. op_name: The logical op that this kernel implements (e.g. ``"rms_norm"``, ``"moe_experts"``). Matches the ``OpSlot``'s ``op_name``. variant: Sub-variant of the op, used when a single op has multiple forward-compatible shapes (e.g. ``"standard"`` vs ``"qwen3_5"`` RMSNorm). Kernels for different variants never collide. factory: Zero-argument callable returning the concrete kernel callable. Kept lazy so optional imports (Liger, Triton, etc.) only load on demand. hardware: Hardware gate enforced at ``resolve()`` time; raises ``RuntimeError`` early when the requested kernel cannot run on the current accelerator. description: Free-form human-readable description, surfaced in registry listings. """ name: str op_name: str variant: str factory: Callable[[], Callable] hardware: HardwareRequirement description: str = "" class KernelRegistry: """Global registry mapping (op_name, variant) -> {impl_name: KernelSpec}.""" def __init__(self): self._specs: dict[tuple[str, str], dict[str, KernelSpec]] = {} def register(self, spec: KernelSpec, force=False) -> None: key = (spec.op_name, spec.variant) bucket = self._specs.setdefault(key, {}) if spec.name in bucket: if force: logger.info( f"Kernel(op='{spec.op_name}', variant='{spec.variant}', name='{spec.name}') is replaced with a new one from {spec.factory.__code__.co_filename}" ) else: raise ValueError( f"Duplicate kernel registration: op='{spec.op_name}', variant='{spec.variant}', name='{spec.name}'" ) bucket[spec.name] = spec def resolve(self, op_name: str, variant: str, impl_name: str) -> Callable | None: """Resolve an implementation by name. Returns ``None`` when *impl_name* is ``"eager"`` (meaning: use the original HF code path). Raises ``KeyError`` if *impl_name* is unknown, and ``RuntimeError`` if the hardware requirement is not satisfied. """ if impl_name == "eager": return None key = (op_name, variant) bucket = self._specs.get(key, {}) if impl_name not in bucket: available = list(bucket.keys()) + ["eager"] raise KeyError( f"Unknown kernel '{impl_name}' for op='{op_name}', variant='{variant}'. Available: {available}" ) spec = bucket[impl_name] if not spec.hardware.is_satisfied(): cc_min = spec.hardware.min_compute_capability cc_max = spec.hardware.max_compute_capability if cc_min is not None and cc_max is not None and cc_min == cc_max: cc_clause = f", compute_capability=={cc_min}" elif cc_min is not None and cc_max is not None: cc_clause = f", {cc_min}<=compute_capability<={cc_max}" elif cc_min is not None: cc_clause = f", compute_capability>={cc_min}" elif cc_max is not None: cc_clause = f", compute_capability<={cc_max}" else: cc_clause = "" raise RuntimeError( f"Kernel '{impl_name}' for op='{op_name}' requires " f"device_type='{spec.hardware.device_type}'" + cc_clause + ", but the current hardware does not satisfy this." ) return spec.factory() def list_available(self, op_name: str, variant: str) -> list[str]: key = (op_name, variant) return list(self._specs.get(key, {}).keys()) KERNEL_REGISTRY = KernelRegistry()