sgl-project--sglang/python/sglang/srt/hardware_backend/mlx/aot.py

"""AOT kernel selection and decode-context helpers for the MLX backend."""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Optional

import mlx.core as mx

from sglang.srt.environ import envs

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import (
        ContiguousAttentionKVCache,
    )


def _load_metal_rope_pool_fused():
    try:
        from sgl_kernel import metal
    except ImportError as exc:
        raise ImportError(
            "sgl_kernel.metal is not importable. Install sgl-kernel in the "
            "active environment before enabling SGLANG_MLX_USE_CUSTOM_ROPE."
        ) from exc

    import_error = getattr(metal, "_IMPORT_ERROR", None)
    if getattr(metal, "_metal", None) is None or import_error is not None:
        reason = f" Reason: {import_error}." if import_error is not None else ""
        raise ImportError(
            "sgl_kernel.metal is importable, but the native Metal extension "
            f"or metallib is not available.{reason} Install the Metal kernels "
            "with `uv run sgl-kernel/setup_metal.py install` from the SGLang "
            "repo root in the active environment."
        ) from import_error
    return metal.rope_pool_fused


@dataclass
class MlxAOTRoPEKernel:
    base: float = 0.0
    config: dict[str, Any] = field(default_factory=dict)
    rope_pool_fused: Optional[Any] = None

    @property
    def enabled(self) -> bool:
        return (
            self.base > 0.0 and bool(self.config) and self.rope_pool_fused is not None
        )


@dataclass
class MlxAOTKernelBuildInputs:
    sample_attn: Any
    n_kv_heads: int
    head_dim: int


@dataclass(frozen=True)
class MlxAOTKernelSpec:
    name: str
    kernel_attr: str
    is_enabled: Callable[[], bool]
    build: Callable[[MlxAOTKernelBuildInputs], Any]


@dataclass
class MlxAOTKernelSet:
    rope: MlxAOTRoPEKernel = field(default_factory=MlxAOTRoPEKernel)
    selected_kernel_names: tuple[str, ...] = ()


class MlxAOTKernelRegistry:
    """Registry for optional MLX AOT kernels.

    Each spec owns one kernel field on ``MlxAOTKernelSet``. The registry is the
    only place that checks kernel enablement policy and model support.
    """

    def __init__(self, specs: tuple[MlxAOTKernelSpec, ...]):
        self._specs = specs

    @property
    def registered_kernel_names(self) -> tuple[str, ...]:
        return tuple(spec.name for spec in self._specs)

    def build_kernel_set(
        self,
        *,
        sample_attn: Any,
        n_kv_heads: int,
        head_dim: int,
    ) -> MlxAOTKernelSet:
        inputs = MlxAOTKernelBuildInputs(
            sample_attn=sample_attn,
            n_kv_heads=n_kv_heads,
            head_dim=head_dim,
        )
        kernel_set = MlxAOTKernelSet()
        selected_kernel_names = []
        for spec in self._specs:
            if not spec.is_enabled():
                continue
            kernel = spec.build(inputs)
            if getattr(kernel, "enabled", False):
                if not hasattr(kernel_set, spec.kernel_attr):
                    raise ValueError(
                        f"AOT kernel {spec.name} targets unknown kernel-set "
                        f"attribute {spec.kernel_attr}"
                    )
                setattr(kernel_set, spec.kernel_attr, kernel)
                selected_kernel_names.append(spec.name)
        kernel_set.selected_kernel_names = tuple(selected_kernel_names)
        if kernel_set.selected_kernel_names:
            logger.info(
                "MLX AOT kernels selected: %s",
                ", ".join(kernel_set.selected_kernel_names),
            )
        return kernel_set


def _build_rope_kernel(inputs: MlxAOTKernelBuildInputs) -> MlxAOTRoPEKernel:
    from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import (
        get_num_heads,
    )

    sample_attn = getattr(inputs.sample_attn, "_inner", inputs.sample_attn)
    rope = getattr(sample_attn, "rope", None)
    if rope is None or getattr(rope, "traditional", False):
        return MlxAOTRoPEKernel()

    rope_dim = int(getattr(rope, "dims", 0))
    if rope_dim == 0:
        return MlxAOTRoPEKernel()
    if rope_dim != inputs.head_dim:
        # AOT kernel currently requires rope_dim == head_dim.
        return MlxAOTRoPEKernel()

    base = float(getattr(rope, "base", 10000.0))
    num_qo_heads = get_num_heads(sample_attn)
    if num_qo_heads is None:
        return MlxAOTRoPEKernel()
    config = {
        "head_dim": int(inputs.head_dim),
        "rope_dim": rope_dim,
        "num_qo_heads": int(num_qo_heads),
        "num_kv_heads": int(inputs.n_kv_heads),
    }
    try:
        rope_pool_fused = _load_metal_rope_pool_fused()
    except Exception as exc:  # noqa: BLE001
        logger.info(
            "AOT Metal RoPE kernel not available (%s) - falling back to "
            "mx.fast.rope.",
            exc,
        )
        return MlxAOTRoPEKernel()

    logger.info(
        f"AOT Metal RoPE kernel ENABLED: head_dim={inputs.head_dim}, "
        f"n_heads={config['num_qo_heads']}, n_kv={config['num_kv_heads']}, "
        f"base={base}"
    )
    return MlxAOTRoPEKernel(
        base=base,
        config=config,
        rope_pool_fused=rope_pool_fused,
    )


MLX_AOT_KERNEL_REGISTRY = MlxAOTKernelRegistry(
    specs=(
        MlxAOTKernelSpec(
            name="metal_rope_pool_fused",
            kernel_attr="rope",
            is_enabled=lambda: envs.SGLANG_MLX_USE_CUSTOM_ROPE.get(),
            build=_build_rope_kernel,
        ),
    )
)


@dataclass
class MlxAOTRoPEContext:
    kernel: MlxAOTRoPEKernel
    kv_pool: Any
    new_token_slots: Optional[mx.array] = None


@dataclass
class MlxAOTKernelContext:
    rope: Optional[MlxAOTRoPEContext] = None

    @classmethod
    def from_decode(
        cls,
        *,
        aot_kernels: MlxAOTKernelSet,
        kv_pool: Any | None,
        req_ids: list[str],
        req_pool_idx: dict[str, int],
        req_to_token_pool: Any | None,
        layer_caches: list[list[ContiguousAttentionKVCache]],
    ) -> MlxAOTKernelContext:
        """Build optional AOT context for one batched decode step."""
        if not aot_kernels.rope.enabled or kv_pool is None:
            return cls()

        new_token_slots = None
        if req_to_token_pool is not None:
            try:
                slot_ids = []
                for req_idx, req_id in enumerate(req_ids):
                    pool_idx = req_pool_idx.get(req_id)
                    if pool_idx is None:
                        raise KeyError(req_id)
                    slot = int(
                        req_to_token_pool.req_to_token[
                            pool_idx, layer_caches[0][req_idx].offset
                        ].item()
                    )
                    slot_ids.append(slot)
                new_token_slots = mx.array(slot_ids, dtype=mx.int32)
            except Exception as exc:  # noqa: BLE001
                logger.warning(
                    "AOT RoPE: failed to resolve new-token slots (%s); "
                    "falling back to RoPE-only for this decode step",
                    exc,
                )

        return cls(
            rope=MlxAOTRoPEContext(
                kernel=aot_kernels.rope,
                kv_pool=kv_pool,
                new_token_slots=new_token_slots,
            )
        )