# 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 logging import os from contextlib import contextmanager from dataclasses import dataclass, field from enum import Enum from pathlib import Path from typing import Any, Callable, Generator from tokenspeed_kernel.platform import PlatformInfo, current_platform from tokenspeed_kernel.registry import KernelRegistry, KernelSpec from tokenspeed_kernel.signature import FormatSignature logger = logging.getLogger(__name__) __all__ = [ "NoKernelFoundError", "SelectedKernel", "SelectionObjective", "SelectionStrategy", "ScoreBreakdown", "SelectionOracle", "AutotuneParams", "SelectionPolicy", "select_kernel", "set_selection_policy", "register_oracle", "kernel_override", "load_config_overrides", "clear_config_overrides", "explain_selection", "spec_matches_traits", "ref_compatible_with_spec", "spec_matches_shape_traits", "warmup_selection", ] class NoKernelFoundError(RuntimeError): """Raised when no kernel matches the requested operation.""" pass class SelectedKernel: """Result of kernel selection — a callable that also carries the kernel name.""" __slots__ = ("name", "impl") def __init__(self, name: str, impl: Callable) -> None: self.name = name self.impl = impl def __call__(self, *args: Any, **kwargs: Any) -> Any: return self.impl(*args, **kwargs) def __repr__(self) -> str: return f"SelectedKernel(name={self.name!r})" class SelectionObjective(Enum): """Objectives are a closed enum — they directly control scoring logic, so an unknown value would silently fall through with no effect. """ DEFAULT = "default" # Balanced heuristic (priority-weighted) LATENCY = "latency" # Minimize per-call latency THROUGHPUT = "throughput" # Maximize tokens/second for large batches PORTABILITY = "portability" # Prefer solutions that work across vendors (Triton) DETERMINISM = "determinism" # Prefer bit-reproducible implementations DEBUG = "debug" # Prefer readable implementations (reference, Triton) class SelectionStrategy(Enum): HEURISTIC = "heuristic" # Score-based ranking (default, instant) AUTOTUNE = "autotune" # Benchmark candidates, pick fastest @dataclass class ScoreBreakdown: """Per-kernel scoring breakdown across all dimensions. Ranking is lexicographic on ``(oracle, objective, priority)`` — the oracle's per-family knowledge wins first, then objective alignment, with the kernel's declared priority band as the final tiebreaker. """ priority: int # [0, 20) — from KernelSpec.priority objective: int # 0 or 1 — 1 if the kernel matches the requested objective oracle: int # [0, 20) — per-family oracle adjustment def sort_key(self) -> tuple[int, int, int]: """Lex sort key (descending — higher is better).""" return (self.oracle, self.objective, self.priority) def __str__(self) -> str: return f"ora={self.oracle} obj={self.objective} pri={self.priority}" class SelectionOracle: """Base class for per-family selection adjustments. Return a score in [0, 20). 10 = neutral. Higher = better fit. """ def adjust( self, spec: KernelSpec, platform: PlatformInfo, traits: dict[str, Any] | None, ) -> int: return 10 @dataclass class AutotuneParams: """Tuning knobs for autotune strategy.""" warmup_iters: int = 3 bench_iters: int = 10 use_cuda_events: bool = True @dataclass class SelectionPolicy: """Per-op selection strategy configuration.""" # Default strategy for all ops default_strategy: SelectionStrategy = SelectionStrategy.HEURISTIC # Per-op overrides: (family, mode) -> strategy op_strategies: dict[tuple[str, str], SelectionStrategy] = field( default_factory=dict ) # Autotune parameters (used when strategy is AUTOTUNE) autotune_params: AutotuneParams = field(default_factory=AutotuneParams) def get_strategy(self, family: str, mode: str) -> SelectionStrategy: return self.op_strategies.get((family, mode), self.default_strategy) @dataclass class _ConfigOverrideEntry: """A single override entry parsed from overrides.yaml.""" name: str | None = None # Exact kernel name solution: str | None = None # Solution backend to match objective: str | None = None # SelectionObjective value string _policy = SelectionPolicy() _oracles: dict[str, SelectionOracle] = {} _global_overrides: dict[tuple[str, str], str] = {} _config_overrides: dict[tuple[str, str], _ConfigOverrideEntry] | None = None def set_selection_policy(policy: SelectionPolicy) -> None: """Set per-op selection strategy. Clears all cached selections.""" global _policy _policy = policy KernelRegistry.get().clear_cache() def register_oracle(family: str, oracle: SelectionOracle) -> None: """Register a per-family selection oracle.""" _oracles[family] = oracle def _get_oracle(family: str) -> SelectionOracle | None: return _oracles.get(family) def _parse_overrides( raw: dict, ) -> dict[tuple[str, str], _ConfigOverrideEntry]: """Parse the ``overrides`` section of the YAML config. Accepted formats:: overrides: attention.decode: solution: flashinfer gemm.mm: name: gluon_gemm_mm_fp8 moe.experts: objective: determinism norm.rmsnorm: triton_rmsnorm # shorthand: treated as kernel name """ result: dict[tuple[str, str], _ConfigOverrideEntry] = {} if not isinstance(raw, dict): return result for op_key, entry in raw.items(): parts = str(op_key).split(".", 1) if len(parts) != 2: logger.warning("Invalid override key '%s' (expected 'family.mode')", op_key) continue family, mode = parts if isinstance(entry, str): result[(family, mode)] = _ConfigOverrideEntry(name=entry) elif isinstance(entry, dict): name = entry.get("name") solution = entry.get("solution") objective = entry.get("objective") result[(family, mode)] = _ConfigOverrideEntry( name=str(name) if name else None, solution=str(solution) if solution else None, objective=str(objective) if objective else None, ) else: logger.warning("Invalid override value for '%s': %r", op_key, entry) return result def load_config_overrides(path: str | os.PathLike[str] | None = None) -> None: """Load kernel overrides from a YAML config file. Args: path: Path to the YAML file. If *None*, uses the ``TOKENSPEED_KERNEL_OVERRIDES_FILE`` env var or falls back to ``~/.config/tokenspeed-kernel/overrides.yaml``. """ global _config_overrides if path is None: env_path = os.environ.get("TOKENSPEED_KERNEL_OVERRIDES_FILE") if env_path: path = Path(env_path) else: path = Path("~/.config/tokenspeed-kernel/overrides.yaml").expanduser() else: path = Path(path) _config_overrides = {} if not path.exists(): return try: import yaml # type: ignore[import-untyped] except ImportError: logger.warning( "PyYAML not installed; cannot load overrides from %s. " "Install with: pip install pyyaml", path, ) return try: with open(path) as f: data = yaml.safe_load(f) except Exception: logger.warning("Failed to load overrides from %s", path, exc_info=True) return if not isinstance(data, dict): return _config_overrides = _parse_overrides(data.get("overrides", {})) if _config_overrides: logger.debug( "Loaded %d config override(s) from %s", len(_config_overrides), path, ) KernelRegistry.get().clear_cache() def clear_config_overrides() -> None: """Clear loaded config overrides. After this call no config-file overrides are active. Call :func:`load_config_overrides` again to reload from a file. """ global _config_overrides _config_overrides = {} def _get_config_override(family: str, mode: str) -> _ConfigOverrideEntry | None: """Return the config-file override for *(family, mode)*, lazily loading from the default path on first access.""" global _config_overrides if _config_overrides is None: load_config_overrides() return _config_overrides.get((family, mode)) # type: ignore[union-attr] def _make_cache_key( family: str, mode: str, format_signature: FormatSignature, arch: str, objective: SelectionObjective, features: frozenset[str] | None, traits: dict[str, Any] | None, solution: str | None = None, ) -> tuple: """Build a hashable cache key including selection-relevant traits.""" traits_key = tuple(sorted(traits.items())) if traits else () mods_key = frozenset(features) if features else frozenset() return ( family, mode, format_signature, arch, objective, mods_key, traits_key, solution, ) def _score_priority(spec: KernelSpec) -> int: """Priority dimension: kernel's inherent quality/maturity.""" return max(0, min(19, spec.priority)) _OBJECTIVE_TAG: dict[SelectionObjective, str] = { SelectionObjective.LATENCY: "latency", SelectionObjective.THROUGHPUT: "throughput", SelectionObjective.PORTABILITY: "portability", SelectionObjective.DETERMINISM: "determinism", SelectionObjective.DEBUG: "determinism", } def _score_objective(spec: KernelSpec, objective: SelectionObjective) -> int: """Objective dimension: 1 if the kernel declares the matching tag, else 0. DEFAULT returns 0 so every kernel ties on this dimension and ranking falls through to oracle/priority. """ tag = _OBJECTIVE_TAG.get(objective) return 1 if tag is not None and tag in spec.tags else 0 def _score_oracle( spec: KernelSpec, platform: PlatformInfo, traits: dict[str, Any] | None, ) -> int: """Oracle dimension: per-family domain-specific scoring.""" oracle = _get_oracle(spec.family) if oracle is None: return 10 # Neutral when no oracle is registered score = oracle.adjust(spec, platform, traits) return max(0, min(19, score)) def _score( spec: KernelSpec, objective: SelectionObjective, platform: PlatformInfo, traits: dict[str, Any] | None, ) -> ScoreBreakdown: """Score a kernel across all ranking dimensions.""" return ScoreBreakdown( priority=_score_priority(spec), objective=_score_objective(spec, objective), oracle=_score_oracle(spec, platform, traits), ) def _rank_by_objective( specs: list[KernelSpec], objective: SelectionObjective, platform: PlatformInfo, traits: dict[str, Any] | None, ) -> list[tuple[KernelSpec, ScoreBreakdown]]: """Rank kernels lexicographically by (oracle, objective, priority). Higher is better. Oracle wins first because per-family oracles encode the most domain knowledge; objective alignment breaks ties next; the kernel's declared priority band is the final tiebreaker. """ scored = [(spec, _score(spec, objective, platform, traits)) for spec in specs] scored.sort(key=lambda x: x[1].sort_key(), reverse=True) return scored def _trait_value_matches(spec_values: frozenset[Any], trait_value: Any) -> bool: if not isinstance(trait_value, (set, frozenset)): trait_value = frozenset({trait_value}) return trait_value.issubset(spec_values) def _ispp_satisfies_alignment(spec: KernelSpec, ispp: Any) -> bool: alignments = spec.traits.get("ispp_alignment") if alignments is None: return True try: ispp_value = int(ispp) except (TypeError, ValueError): return False return any( int(alignment) > 0 and ispp_value % int(alignment) == 0 for alignment in alignments ) def spec_matches_traits( spec: KernelSpec, traits: dict[str, Any], *, require_all_traits: bool = False, ) -> bool: """Return whether a spec's declared traits match the requested traits. Args: spec: Registered kernel specification to test. traits: Trait requirements. Values may be concrete scalars (for example, ``{"head_dim": 128}``) or sets/frozensets of allowed values. require_all_traits: When ``False`` (selection behavior), unknown traits on the spec are ignored. When ``True`` (reference compatibility checks), every requested trait must be explicitly present on the spec. """ for trait_name, trait_value in traits.items(): # ispp stands for "intermediate size per partition" and has special # alignment requirements that depend on the kernel's declared # supported alignments (if any). It is used in some MoE ops to ensure # the intermediate buffer sizes are compatible with the kernel's # requirements. if trait_name == "ispp": if not _ispp_satisfies_alignment(spec, trait_value): return False continue spec_values = spec.traits.get(trait_name) if spec_values is None: if require_all_traits: return False continue if not _trait_value_matches(spec_values, trait_value): return False return True def ref_compatible_with_spec(ref: KernelSpec, spec: KernelSpec) -> bool: """Return whether a reference kernel can handle the same inputs as a test kernel. For each trait the reference declares, the spec must declare that same trait with values that fully cover the reference's required values. Traits the reference does not declare are unconstrained (the reference is general with respect to those traits). """ for trait_name, ref_values in ref.traits.items(): spec_values = spec.traits.get(trait_name) if spec_values is None: return False if not ref_values.issubset(spec_values): return False return True def spec_matches_shape_traits(spec: KernelSpec, shape: dict[str, Any]) -> bool: """Return whether a spec's alignment traits match a concrete shape.""" alignment_traits: dict[str, tuple[str, int]] = { "n_align_16": ("N", 16), "n_align_64": ("N", 64), "n_align_128": ("N", 128), "k_align_16": ("K", 16), "k_align_64": ("K", 64), "k_align_128": ("K", 128), } for trait_name, (dim_name, alignment) in alignment_traits.items(): values = spec.traits.get(trait_name) if values is None or True not in values: continue dim = shape.get(dim_name) if isinstance(dim, int) and dim % alignment != 0: return False return True def _filter_by_traits( specs: list[KernelSpec], traits: dict[str, Any], ) -> list[KernelSpec]: """Filter kernels by op-specific trait compatibility.""" return [spec for spec in specs if spec_matches_traits(spec, traits)] def _resolve_override( registry: KernelRegistry, family: str, mode: str, format_signature: object, override: str, platform: PlatformInfo, ) -> SelectedKernel: impl = registry.get_impl(override) if impl is not None: return SelectedKernel(name=override, impl=impl) specs = registry.get_for_operator(family, mode, solution=override) if specs: kernel_name = specs[0].name impl = registry.get_impl(kernel_name) if impl is not None: return SelectedKernel(name=kernel_name, impl=impl) raise NoKernelFoundError( f"Override '{override}' not found for {family}.{mode} ({format_signature})" ) def _log_selection( family: str, mode: str, format_signature: object, winner: KernelSpec, scored: list[tuple[KernelSpec, ScoreBreakdown]], platform: PlatformInfo, objective: SelectionObjective, ) -> None: """Log selection result if verbose mode is enabled.""" if not os.environ.get("TOKENSPEED_KERNEL_VERBOSE"): return breakdown = next((s for spec, s in scored if spec.name == winner.name), None) if breakdown: logger.info( "[tokenspeed_kernel] %s.%s(%s) -> %s (%s, %s)", family, mode, format_signature, winner.name, breakdown, platform.arch, ) else: logger.info( "[tokenspeed_kernel] %s.%s(%s) -> %s (%s)", family, mode, format_signature, winner.name, platform.arch, ) def select_kernel( family: str, mode: str, format_signature: FormatSignature, *, features: frozenset[str] | None = None, platform: PlatformInfo | None = None, objective: SelectionObjective = SelectionObjective.DEFAULT, traits: dict[str, Any] | None = None, solution: str | None = None, override: str | None = None, expected_kernel_name: str | None = None, ) -> SelectedKernel: """Select the best kernel for an operation. On first call for a given (family, mode, format_signature, platform, objective, traits, solution) combination, runs the full selection pipeline. Subsequent calls with the same arguments return the cached result — a single dict lookup. Args: family: Operator family (e.g., "attention") mode: Operator mode (e.g., "decode") format_signature: Role-indexed tensor format signature features: Required operator features (e.g., {"paged"}) platform: Hardware to match (auto-detected if None) objective: Selection objective (see SelectionObjective enum) traits: Op-specific trait values that affect kernel applicability (e.g., {"head_dim": 128, "num_kv_heads": 8}) solution: Restrict selection to a registered solution while preserving normal platform, format signature, and trait filtering. override: Force a specific kernel name or solution string expected_kernel_name: Debug-only hint. When set, a warning is logged if the selected kernel differs from this name. The selected kernel is still used regardless of the mismatch. Returns: A :class:`SelectedKernel` that is directly callable and also exposes the winning kernel's ``name``. """ platform = platform or current_platform() # --- Override resolution (lowest → highest priority) --- # 1. Config file (lowest priority) config_entry = _get_config_override(family, mode) if config_entry is not None: if override is None and solution is None: if config_entry.name: override = config_entry.name elif config_entry.solution: solution = config_entry.solution if config_entry.objective and objective == SelectionObjective.DEFAULT: try: objective = SelectionObjective(config_entry.objective) except ValueError: logger.warning( "Invalid objective '%s' in overrides config for %s.%s", config_entry.objective, family, mode, ) # 2. Context-manager global overrides global_override = _global_overrides.get((family, mode)) if global_override: override = global_override # 3. Environment variable env_key = f"TOKENSPEED_KERNEL_OVERRIDE_{family.upper()}_{mode.upper()}" env_override = os.environ.get(env_key) if env_override: override = env_override registry = KernelRegistry.get() # Fast path: check cache (skipped when override is active) cache_key = _make_cache_key( family, mode, format_signature, platform.arch, objective, features, traits, solution, ) if override is None: cached = registry.cache_get(cache_key) if cached is not None: return cached if override: return _resolve_override( registry, family, mode, format_signature, override, platform ) # Get candidates (same filtering for both strategies) candidates = registry.get_for_operator( family, mode, features=features, platform=platform, format_signature=format_signature, solution=solution, ) solution_clause = f" with solution {solution!r}" if solution else "" if not candidates: raise NoKernelFoundError( f"No kernel found for {family}.{mode} ({format_signature})" f"{solution_clause} on {platform.device_name}" ) if traits: candidates = _filter_by_traits(candidates, traits) if not candidates: raise NoKernelFoundError( f"No kernel found for {family}.{mode} ({format_signature})" f"{solution_clause} with traits {traits} on {platform.device_name}" ) # Strategy dispatch strategy = _policy.get_strategy(family, mode) if strategy == SelectionStrategy.AUTOTUNE: winner, scored = _autotune_select( candidates, family, mode, format_signature, platform, traits, _policy.autotune_params, ) else: scored = _rank_by_objective(candidates, objective, platform, traits) winner = scored[0][0] _log_selection(family, mode, format_signature, winner, scored, platform, objective) if expected_kernel_name and winner.name != expected_kernel_name: logger.warning( "[tokenspeed_kernel] select_kernel(%s.%s, %s) chose '%s' but " "expected '%s'. Score breakdown — selected: %s", family, mode, format_signature, winner.name, expected_kernel_name, next((s for sp, s in scored if sp.name == winner.name), "N/A"), ) impl = registry.get_impl(winner.name) result = SelectedKernel(name=winner.name, impl=impl) registry.cache_put(cache_key, result) return result def _autotune_select( candidates: list[KernelSpec], family: str, mode: str, format_signature: object, platform: PlatformInfo, traits: dict[str, Any] | None, params: AutotuneParams, ) -> tuple[KernelSpec, list[tuple[KernelSpec, ScoreBreakdown]]]: """Benchmark candidates and return the fastest. Falls back to heuristic ranking when the autotuning infrastructure (input generators, benchmark runner) is not yet available. """ scored = _rank_by_objective( candidates, SelectionObjective.DEFAULT, platform, traits ) winner = scored[0][0] logger.debug( "[tokenspeed_kernel:autotune] falling back to heuristic for %s.%s(%s)", family, mode, format_signature, ) return winner, scored @contextmanager def kernel_override( family: str, mode: str, kernel_name: str ) -> Generator[None, None, None]: """Context manager for scoped kernel override.""" key = (family, mode) old = _global_overrides.get(key) _global_overrides[key] = kernel_name try: yield finally: if old is None: _global_overrides.pop(key, None) else: _global_overrides[key] = old def explain_selection( family: str, mode: str, format_signature: FormatSignature, *, features: frozenset[str] | None = None, platform: PlatformInfo | None = None, objective: SelectionObjective = SelectionObjective.DEFAULT, traits: dict[str, Any] | None = None, solution: str | None = None, ) -> str: """Return a human-readable explanation of kernel selection. Example output:: Op: attention.decode (bfloat16) Platform: NVIDIA H100 (sm_90) Objective: default Ranking: lex (oracle, objective, priority); higher wins Candidates (3 matched, 5 registered): 1. flashinfer_decode [SELECTED] ora=16 obj=1 pri=14 2. triton_decode ora=10 obj=0 pri=10 Filtered out: - aiter_decode: vendor mismatch (requires amd) """ platform = platform or current_platform() registry = KernelRegistry.get() all_specs = registry.list_kernels(family=family, mode=mode) candidates = registry.get_for_operator( family, mode, features=features, platform=platform, format_signature=format_signature, solution=solution, ) if traits: candidates = _filter_by_traits(candidates, traits) scored = _rank_by_objective(candidates, objective, platform, traits) filtered_names = {s.name for s in candidates} filtered_out = [s for s in all_specs if s.name not in filtered_names] lines = [ f"Op: {family}.{mode} ({format_signature})", f"Platform: {platform.device_name} ({platform.arch})", f"Solution: {solution or 'any'}", f"Objective: {objective.value}", "Ranking: lex (oracle, objective, priority); higher wins", "", f"Candidates ({len(scored)} matched, {len(all_specs)} registered):", ] for i, (spec, breakdown) in enumerate(scored): marker = " [SELECTED]" if i == 0 else "" lines.append(f" {i + 1}. {spec.name}{marker}") lines.append(f" {breakdown}") if filtered_out: lines.append("") lines.append("Filtered out:") for spec in filtered_out: reasons: list[str] = [] if ( spec.capability.vendors and platform.vendor not in spec.capability.vendors ): reasons.append( f"vendor mismatch (requires " f"{', '.join(spec.capability.vendors)})" ) missing = spec.capability.missing_features(platform) if missing: reasons.append(f"missing features: {', '.join(missing)}") if spec.capability.min_arch_version: if not (platform.arch_version >= spec.capability.min_arch_version): reasons.append( f"arch mismatch (requires " f"{spec.capability.min_arch_version})" ) if format_signature and not spec.supports_format_signature( format_signature ): reasons.append( f"format signature mismatch (supports " f"{', '.join(str(d) for d in spec.format_signatures)})" ) if solution and spec.solution != solution: reasons.append(f"solution mismatch (is {spec.solution!r})") reason_str = "; ".join(reasons) if reasons else "unknown" lines.append(f" - {spec.name}: {reason_str}") return "\n".join(lines) def warmup_selection( ops: list[tuple[str, str, FormatSignature, dict | None]] | None = None, ) -> None: """Pre-resolve kernel selection for explicit op signatures. Pass ``ops`` from model initialization to front-load heuristic and autotune costs for the actual hot-path call sites. Each entry must include the exact ``FormatSignature`` and trait values used by runtime selection. When ``ops`` is ``None``, this performs only a deterministic smoke warmup: one representative signature for each registered operator, with no traits. That path verifies the registry and fills a small cache sample, but it does not warm all supported signatures, trait combinations, or feature-specific call paths. """ if ops is None: ops = [] registry = KernelRegistry.get() for family, mode in registry.list_operators(): specs = registry.get_for_operator(family, mode) if not specs or not specs[0].format_signatures: continue # No-arg warmup is intentionally a smoke path. Pick a stable # representative from the highest-priority spec; callers that need # comprehensive warmup should pass explicit op signatures. format_signature = sorted(specs[0].format_signatures, key=str)[0] ops.append((family, mode, format_signature, None)) for family, mode, format_signature, traits in ops: try: select_kernel(family, mode, format_signature, traits=traits) except NoKernelFoundError: logger.debug( "[tokenspeed_kernel] warmup: no kernel for %s.%s(%s)", family, mode, format_signature, )