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

"""MLX model runner for Apple Silicon.

Slot allocation and radix-trie prefix matching are handled by the
scheduler (``TokenToKVPoolAllocator`` / ``RadixCache``).  This runner
reads cached attention KV from ``MlxAttentionKVPool``, restores any
native auxiliary layer state, runs the forward pass, and writes the new
cache state back.  Each request keeps model-shaped cache entries:
attention layers use ``ContiguousAttentionKVCache`` and auxiliary layers
use native ``mlx-lm`` cache objects.

The module also exposes a lazy-eval (`*_start` / `*_finalize`) surface
used by the MLX overlap scheduler to pipeline CPU bookkeeping with
GPU execution.  The lazy API is a thin split of the synchronous API:
``*_start`` builds the compute graph without materialising outputs,
``*_finalize`` blocks on the lazy token(s) and commits per-request
state.
"""

import logging
import time
from dataclasses import dataclass
from typing import Any

import mlx.core as mx
import psutil
from mlx.utils import tree_flatten
from mlx_lm import load as mlx_lm_load
from mlx_lm.utils import quantize_model as mlx_lm_quantize_model

from sglang.srt.environ import envs
from sglang.srt.hardware_backend.mlx.aot import (
    MLX_AOT_KERNEL_REGISTRY,
    MlxAOTKernelSet,
)
from sglang.srt.hardware_backend.mlx.kv_cache import (
    AttentionOffsetCache,
    BatchedDecodeContext,
    ContiguousAttentionKVCache,
    MlxAttentionKVPool,
    MLXAttentionWrapper,
    MlxModelCacheLayout,
    PoolBackedAttentionKVCache,
    clear_context,
    find_attention_layers,
    get_head_dim,
    get_num_kv_heads,
    patch_model_attention,
    set_context,
    uses_sliding_window_attention,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.runtime_context import get_server_args

logger = logging.getLogger(__name__)


@dataclass
class MlxPendingPrefill:
    """Lazy prefill state, finalised after ``mx.eval``/``async_eval``.

    ``cache`` is the per-layer cache list that will
    become ``_req_caches[req_id]`` once the request is committed.  It
    may have been converted from transient pool-backed attention caches
    already, so its ``state`` arrays are safe to hand to ``async_eval``.
    """

    lazy_token: mx.array
    cache: list[Any]
    req_id: str
    full_token_ids: list[int]
    req_pool_idx: int
    synced_offset: int


@dataclass
class MlxPendingExtend:
    """Lazy chunked-prefill-continuation state for an existing request.

    Mirrors :meth:`MlxModelRunner.extend` split into launch/finalize
    halves.  ``cache`` is the request's existing per-layer cache (not a
    fresh one) so the graph writes extend onto the already-materialised
    prefix.
    """

    lazy_token: mx.array
    req_id: str
    new_token_ids: list[int]
    new_synced_offset: int


@dataclass
class MlxPendingDecode:
    """Lazy decode state, finalised after ``mx.eval``/``async_eval``.

    ``caches`` is a per-request list of per-layer cache
    references (``caches[req_idx][layer_idx]``).  These are the same
    objects the attention wrapper writes into during the forward pass,
    so :meth:`decode_batch_start_chained` can launch the next step on
    top of the same caches without materialising this step first.
    """

    lazy_tokens: mx.array
    req_ids: list[str]
    caches: list[list[Any]]


_MLX_QUANTIZATION_PRESETS: dict[str, tuple[int, int]] = {
    # name -> (bits, group_size). group_size=64 matches the mlx-community convention.
    "mlx_q4": (4, 64),
    "mlx_q8": (8, 64),
}
_MLX_KV_FLOAT_DTYPES = {mx.float16, mx.bfloat16, mx.float32}


class MlxModelRunner:
    """MLX model runner with radix-cache prefix sharing."""

    def __init__(
        self,
        model_path: str,
        trust_remote_code: bool = False,
        disable_radix_cache: bool = False,
        pool_size: int | None = None,
        mem_fraction_static: float = 0.8,
        quantization: str | None = None,
    ):
        self.model_path = model_path
        self.trust_remote_code = trust_remote_code
        self.model = None
        self.disable_radix_cache = disable_radix_cache
        self._mem_fraction_static = mem_fraction_static
        # Counter used to trigger periodic mx.clear_cache() calls.
        self._decode_step_ct: int = 0
        self._clear_steps = envs.SGLANG_MLX_CLEAR_CACHE_STEPS.get()
        # On-the-fly quantization preset (e.g. "mlx_q4"). None = no on-load quantization.
        # Pre-quantized HF repos load correctly regardless of this setting:
        # mlx_lm.load() detects the config and instantiates QuantizedLinear
        # modules directly.
        self._quantization: str | None = quantization

        self._load_model()

        # Pin MLX allocations to prevent OS paging
        device_info = mx.device_info()
        max_wired = int(device_info.get("max_recommended_working_set_size", 0))
        if max_wired > 0:
            mx.set_wired_limit(max_wired)
            logger.info(f"Wired memory limit set to {max_wired / (1024**3):.1f} GB")

        patch_model_attention(self.model)

        layer_list, attn_attrs = find_attention_layers(self.model)
        self._cache_layout = MlxModelCacheLayout.from_attention_discovery(
            layer_list,
            attn_attrs,
        )
        if self._cache_layout.num_attention_layers == 0:
            raise RuntimeError("MLX model has no supported attention layers")
        if self._cache_layout.has_auxiliary_state and not hasattr(
            self.model, "make_cache"
        ):
            raise RuntimeError(
                "MLX models with auxiliary cache state require model.make_cache()."
            )
        if self._cache_layout.has_auxiliary_state:
            self._model_embed, self._model_norm, self._model_lm_head = (
                self._extract_model_components()
            )
        self._max_seq_len = 4096  # doubles on overflow

        self._req_caches: dict[str, list[Any]] = {}
        self._req_token_ids: dict[str, list[int]] = {}
        self._cache_pool: list[list[Any]] = []  # reusable full-attention caches

        self._attention_kv_pool: MlxAttentionKVPool | None = None
        self._req_to_token_pool: ReqToTokenPool | None = None
        self._req_pool_idx: dict[str, int] = {}
        self._req_synced_offset: dict[str, int] = {}

        self._pool_size = self._compute_pool_size(pool_size)
        self._aot_kernels = self._build_aot_kernels()

    @staticmethod
    def _extract_logits(model_output):
        """Extract logits from model output, handling both tuple and direct returns."""
        if isinstance(model_output, tuple):
            return model_output[0]
        return model_output

    def _new_cache_skeleton(self) -> list[Any]:
        """Create a model-shaped cache list before attention cache wiring."""
        if self._cache_layout.has_auxiliary_state:
            cache = self.model.make_cache()
            if len(cache) != self._cache_layout.num_layers:
                raise RuntimeError(
                    "model.make_cache() returned "
                    f"{len(cache)} entries for {self._cache_layout.num_layers} layers"
                )
        else:
            cache = [None] * self._cache_layout.num_layers
        return cache

    def _new_native_cache(self) -> list[Any]:
        """Create a model-shaped cache list with attention KV adapters."""
        cache = self._new_cache_skeleton()
        for layer_idx in self._cache_layout.attention_layer_indices:
            cache[layer_idx] = ContiguousAttentionKVCache(max_seq_len=self._max_seq_len)
        return cache

    def _acquire_cache(self) -> list[Any]:
        """Get a reusable cache list from the pool, or create a new one."""
        if not self._cache_layout.has_auxiliary_state and self._cache_pool:
            cache = self._cache_pool.pop()
            for c in cache:
                c.offset = 0
            return cache
        return self._new_native_cache()

    def _release_cache(self, cache: list[Any]) -> None:
        """Return a cache list to the pool for reuse."""
        if not self._cache_layout.has_auxiliary_state:
            self._cache_pool.append(cache)

    def _first_attention_cache(self, cache: list[Any]) -> Any:
        return cache[self._cache_layout.first_attention_layer_index]

    def _get_auxiliary_state_pool_index(self, req_pool_idx: int) -> Any | None:
        if (
            not self._cache_layout.has_auxiliary_state
            or self._req_to_token_pool is None
            or not hasattr(self._req_to_token_pool, "get_auxiliary_state_indices")
        ):
            return None
        return self._req_to_token_pool.get_auxiliary_state_indices(req_pool_idx)

    def _get_auxiliary_state_pool(self) -> Any | None:
        return getattr(self._req_to_token_pool, "auxiliary_state_pool", None)

    def _restore_auxiliary_state(self, req_pool_idx: int, cache: list[Any]) -> bool:
        pool_index = self._get_auxiliary_state_pool_index(req_pool_idx)
        pool = self._get_auxiliary_state_pool()
        if pool_index is None or not hasattr(pool, "restore_cache"):
            return False
        return pool.restore_cache(
            pool_index,
            cache,
            self._cache_layout.auxiliary_layer_indices,
        )

    def _store_auxiliary_state(self, req_pool_idx: int, cache: list[Any]) -> None:
        pool_index = self._get_auxiliary_state_pool_index(req_pool_idx)
        pool = self._get_auxiliary_state_pool()
        if pool_index is None or not hasattr(pool, "store_cache"):
            return
        pool.store_cache(
            pool_index,
            cache,
            self._cache_layout.auxiliary_layer_indices,
        )

    def store_auxiliary_state_for_request(self, req_id: str) -> None:
        """Snapshot native auxiliary state before scheduler-owned radix insert."""
        req_pool_idx = self._req_pool_idx.get(req_id)
        cache = self._req_caches.get(req_id)
        if req_pool_idx is None or cache is None:
            return
        self._store_auxiliary_state(req_pool_idx, cache)

    def _select_auxiliary_state_track_len(
        self,
        *,
        prefix_len: int,
        new_token_count: int,
        full_len: int,
        req: Any | None,
    ) -> int | None:
        if (
            not self._cache_layout.has_auxiliary_state
            or req is None
            or new_token_count <= 0
        ):
            return None

        chunk_size = get_server_args().mamba_cache_chunk_size
        track_len = prefix_len + (new_token_count // chunk_size) * chunk_size
        branching_len = getattr(req, "mamba_branching_seqlen", None)
        if (
            branching_len is not None
            and prefix_len < branching_len <= prefix_len + new_token_count
            and (branching_len - prefix_len) % chunk_size == 0
        ):
            track_len = branching_len

        if track_len <= prefix_len or track_len > full_len:
            return None
        return track_len

    def _store_tracked_auxiliary_state(
        self,
        req: Any | None,
        cache: list[Any],
        track_len: int | None,
    ) -> None:
        if (
            req is None
            or track_len is None
            or not self._cache_layout.has_auxiliary_state
        ):
            return
        pool = self._get_auxiliary_state_pool()
        if pool is None or not hasattr(pool, "store_cache"):
            return

        track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None)
        if track_buffer is None:
            track_buffer = pool.alloc(1)
            if track_buffer is None:
                logger.warning(
                    "MLX auxiliary-state track slot allocation failed; "
                    "falling back to leaf-only auxiliary-state radix caching."
                )
                return
            req.mamba_ping_pong_track_buffer = track_buffer
            req.mamba_next_track_idx = 0

        pool.store_cache(
            track_buffer[0],
            cache,
            self._cache_layout.auxiliary_layer_indices,
        )
        req.mamba_last_track_seqlen = track_len

    def _cache_with_pool_backed_attention(
        self, prefix_slot_ids: list[int], prefix_len: int
    ) -> list[Any]:
        assert self._attention_kv_pool is not None
        slot_ids_mx = mx.array(prefix_slot_ids, dtype=mx.int32)
        cache = self._new_cache_skeleton()
        for layer_idx in self._cache_layout.attention_layer_indices:
            cache[layer_idx] = PoolBackedAttentionKVCache(
                self._attention_kv_pool,
                self._cache_layout.attention_pool_index(layer_idx),
                slot_ids_mx,
                prefix_len,
            )
        return cache

    def _materialize_pool_backed_attention(self, cache: list[Any]) -> list[Any]:
        contiguous_cache = self._acquire_cache()
        for layer_idx in self._cache_layout.attention_layer_indices:
            pbc = cache[layer_idx]
            contiguous_cache[layer_idx].update_and_fetch(
                pbc._full_keys, pbc._full_values
            )
        for layer_idx in self._cache_layout.auxiliary_layer_indices:
            contiguous_cache[layer_idx] = cache[layer_idx]
        return contiguous_cache

    @staticmethod
    def _cache_arrays(cache: Any) -> list[mx.array]:
        """Return every MLX array nested under ``cache.state``."""
        arrays: list[mx.array] = []

        def collect(value: Any) -> None:
            if isinstance(value, mx.array):
                arrays.append(value)
            elif value is None:
                return
            elif isinstance(value, (list, tuple)):
                for item in value:
                    collect(item)
            elif isinstance(value, dict):
                for item in value.values():
                    collect(item)

        collect(getattr(cache, "state", ()))
        return arrays

    @staticmethod
    def _eval_with_cache(token_result: mx.array, cache: list[Any]) -> None:
        """Evaluate token result and all cache buffers in one mx.eval call."""
        mx.eval(
            token_result,
            *[s for c in cache for s in MlxModelRunner._cache_arrays(c)],
        )

    @staticmethod
    def _cache_state_arrays(pending_caches: list[list[Any]]) -> list[mx.array]:
        """Flatten pending decode cache state list into an array list.

        Safe to hand to ``mx.async_eval``.
        """
        return [
            s
            for cache_list in pending_caches
            for cache in cache_list
            for s in MlxModelRunner._cache_arrays(cache)
        ]

    def _load_model(self):
        """Load model using mlx_lm. If ``self._quantization`` requests a preset
        (e.g. ``mlx_q4``), quantize fp16 weights in-place via
        :func:`mlx_lm.utils.quantize_model` after load.
        """
        logger.info(f"Loading MLX model: {self.model_path}")
        start_time = time.time()

        # We need the config dict to pass into quantize_model so it knows tied/embedding
        # layout. return_config=True is cheap and ignored when no quantization is requested.
        loaded = mlx_lm_load(
            self.model_path,
            tokenizer_config={"trust_remote_code": self.trust_remote_code},
            return_config=True,
        )
        self.model, _tokenizer, config = loaded

        if self._quantization in _MLX_QUANTIZATION_PRESETS:
            bits, group_size = _MLX_QUANTIZATION_PRESETS[self._quantization]
            # Skip if the model was already loaded quantized (pre-quantized HF repo);
            # mlx_lm.load detects the config and instantiates QuantizedLinear directly,
            # so applying the preset on top would be redundant.
            if "quantization" in (config or {}):
                logger.info(
                    "MLX model is already quantized by the HF repo; "
                    f"ignoring --quantization={self._quantization}"
                )
            else:
                # Read weight-tensor totals from MLX array metadata (shape + dtype).
                # This is zero-cost — neither materializes the lazy fp16 weights nor
                # forces them to be peak-resident in memory at once (which on a 64 GB
                # Mac running a 32 B model would put us within a few GB of OOM).
                bytes_before = sum(
                    p.size * p.itemsize
                    for _, p in tree_flatten(self.model.parameters())
                )
                q_start = time.time()
                logger.info(
                    f"Quantizing MLX model on-the-fly: bits={bits} "
                    f"group_size={group_size} (preset={self._quantization})"
                )
                self.model, _new_config = mlx_lm_quantize_model(
                    self.model,
                    config or {},
                    group_size=group_size,
                    bits=bits,
                )
                bytes_after = sum(
                    p.size * p.itemsize
                    for _, p in tree_flatten(self.model.parameters())
                )
                q_time = time.time() - q_start
                pct_reduction = (1 - bytes_after / max(bytes_before, 1)) * 100
                logger.info(
                    f"Quantization complete in {q_time:.2f}s — "
                    f"weight bytes: {bytes_before / 1024**3:.2f} GB -> "
                    f"{bytes_after / 1024**3:.2f} GB ({pct_reduction:.1f}% reduction)"
                )

        # Force-evaluate weights so mx.get_active_memory() reflects
        # actual usage before attention KV pool sizing.
        mx.eval(self.model.parameters())

        load_time = time.time() - start_time
        logger.info(f"MLX model loaded in {load_time:.2f}s")

        # Optional: Path B fusion — keep up_proj/gate_proj weights separate
        # (no matmul-kernel tile regression) but fuse the swiglu activation
        # into the gate matmul via a custom Metal kernel. Activated by
        # SGLANG_MLX_FUSE_SWIGLU=1. Mutually exclusive with FUSE_SWITCHGLU.
        # See: python/sglang/srt/hardware_backend/mlx/moe/fused_swiglu.py
        if envs.SGLANG_MLX_FUSE_SWIGLU.get():
            from sglang.srt.hardware_backend.mlx.moe.fused_swiglu import (
                patch_switch_glu_with_fused_swiglu,
            )

            n_patched = patch_switch_glu_with_fused_swiglu(self.model)
            logger.info(
                f"MLX SwiGLU activation fusion enabled: patched {n_patched} blocks"
            )

    def _attention_module_for_layer(self, layer_idx: int) -> Any:
        attn = getattr(
            self._cache_layout.layers[layer_idx],
            self._cache_layout.attention_attr(layer_idx),
        )
        if isinstance(attn, MLXAttentionWrapper):
            return attn._inner
        return attn

    def _attention_kv_config_for_layer(
        self, layer_idx: int
    ) -> tuple[int, int, mx.Dtype]:
        layer = self._cache_layout.layers[layer_idx]
        sample_attn = self._attention_module_for_layer(layer_idx)
        if uses_sliding_window_attention(layer, sample_attn):
            raise NotImplementedError(
                "MLX radix attention KV pool does not support sliding-window "
                f"attention yet at layer {layer_idx}. Sliding-window KV needs "
                "per-layer/window-aware pools."
            )
        n_kv_heads = get_num_kv_heads(sample_attn)
        if n_kv_heads is None:
            raise RuntimeError(
                f"Cannot determine n_kv_heads from attention module at layer {layer_idx}"
            )
        head_dim = get_head_dim(sample_attn)
        if head_dim is None:
            raise RuntimeError(
                f"Cannot determine head_dim from attention module at layer {layer_idx}"
            )
        dtype = mx.float16
        if hasattr(sample_attn, "k_proj") and hasattr(sample_attn.k_proj, "weight"):
            dtype = sample_attn.k_proj.weight.dtype
        if dtype not in _MLX_KV_FLOAT_DTYPES:
            # QuantizedLinear packs weights as integers, but the KV cache
            # stores dequantized projection outputs, which are produced in
            # the compute dtype carried by the quantization scales.  Storing
            # at that dtype instead of float32 halves pool bytes per slot
            # and keeps prefix-hit forwards in the same dtype as the no-hit
            # path (a float32 pool promoted every post-hit concat).
            scales = getattr(sample_attn.k_proj, "scales", None)
            if scales is not None and scales.dtype in _MLX_KV_FLOAT_DTYPES:
                dtype = scales.dtype
            else:
                dtype = mx.float32
        return n_kv_heads, head_dim, dtype

    def _get_attn_config(self) -> tuple[int, int, mx.Dtype]:
        """Return the uniform attention KV config used by the shared MLX pool."""
        if self._cache_layout.num_attention_layers == 0:
            raise RuntimeError(
                "Cannot determine attention config: no attention module found"
            )
        first_layer_idx = self._cache_layout.first_attention_layer_index
        first_config = self._attention_kv_config_for_layer(first_layer_idx)
        for layer_idx in self._cache_layout.attention_layer_indices[1:]:
            config = self._attention_kv_config_for_layer(layer_idx)
            if config != first_config:
                raise NotImplementedError(
                    "MLX radix attention KV pool requires uniform softmax-attention "
                    "KV shape across layers. "
                    f"Layer {first_layer_idx} has {first_config}, "
                    f"but layer {layer_idx} has {config}. "
                    "Heterogeneous attention KV or sliding-window KV needs "
                    "per-layer pools."
                )
        return first_config

    def _compute_pool_size(self, explicit_size: int | None) -> int:
        """Determine pool slot count (auto-size from available memory if needed)."""
        if explicit_size is not None:
            return explicit_size
        n_kv_heads, head_dim, dtype = self._get_attn_config()
        num_layers = self._cache_layout.num_attention_layers
        sys_available = psutil.virtual_memory().available
        mlx_limit = mx.device_info().get(
            "max_recommended_working_set_size",
            mx.device_info().get("memory_size", 0),
        )
        mlx_used = mx.get_active_memory()
        mlx_usable = int(mlx_limit * self._mem_fraction_static)
        kv_budget = min(
            max(mlx_usable - mlx_used, 0),
            int(sys_available * self._mem_fraction_static),
        )
        bytes_per_slot = 2 * num_layers * n_kv_heads * head_dim * dtype.size
        pool_size = max(kv_budget // bytes_per_slot, 256)
        logger.info(
            f"Auto-sized attention KV pool: "
            f"sys_available={sys_available / (1024**3):.2f} GB, "
            f"mlx_limit={mlx_limit / (1024**3):.1f} GB, "
            f"mlx_used={mlx_used / (1024**3):.2f} GB, "
            f"kv_budget={kv_budget / (1024**3):.2f} GB, "
            f"bytes_per_slot={bytes_per_slot}, pool_size={pool_size}"
        )
        return pool_size

    @property
    def pool_size(self) -> int:
        return self._pool_size

    def _build_aot_kernels(self) -> MlxAOTKernelSet:
        """Build model-level set of optional registered AOT kernels."""
        if self._cache_layout.num_attention_layers == 0:
            return MlxAOTKernelSet()
        layer_idx = self._cache_layout.first_attention_layer_index
        sample_attn = getattr(
            self._cache_layout.layers[layer_idx],
            self._cache_layout.attention_attr(layer_idx),
        )
        n_kv_heads, head_dim, _ = self._get_attn_config()
        return MLX_AOT_KERNEL_REGISTRY.build_kernel_set(
            sample_attn=sample_attn,
            n_kv_heads=int(n_kv_heads),
            head_dim=int(head_dim),
        )

    def init_cache_pools(self, req_to_token_pool: ReqToTokenPool | None) -> None:
        """Create attention KV pool (+1 for padding slot 0)."""
        self._req_to_token_pool = req_to_token_pool
        if self.disable_radix_cache:
            return
        n_kv_heads, head_dim, dtype = self._get_attn_config()
        # +1 for padding slot 0
        self._attention_kv_pool = MlxAttentionKVPool(
            pool_size=self._pool_size + 1,
            num_layers=self._cache_layout.num_attention_layers,
            n_kv_heads=n_kv_heads,
            head_dim=head_dim,
            dtype=dtype,
        )
        logger.info(
            f"Attention KV pool initialized: pool_size={self._pool_size} "
            f"(buffer size {self._pool_size + 1} incl. padding slot 0), "
            f"{self._cache_layout.num_attention_layers} attention layers, "
            f"{n_kv_heads} kv_heads, {head_dim} head_dim"
        )

    def prefill(
        self,
        req_id: str,
        new_token_ids: list[int],
        full_token_ids: list[int],
        prefix_slot_ids: list[int],
        new_slot_ids: list[int],
        req_pool_idx: int,
        req: Any | None = None,
    ) -> int:
        """Prefill a request.  Returns next_token_id."""
        pending = self.prefill_start(
            req_id=req_id,
            new_token_ids=new_token_ids,
            full_token_ids=full_token_ids,
            prefix_slot_ids=prefix_slot_ids,
            new_slot_ids=new_slot_ids,
            req_pool_idx=req_pool_idx,
            req=req,
        )
        self._eval_with_cache(pending.lazy_token, pending.cache)
        return self.prefill_finalize(pending)

    def extend(
        self,
        req_id: str,
        new_token_ids: list[int],
        new_slot_ids: list[int],
    ) -> int:
        """Continue prefill for a chunked request.  Returns next_token_id."""
        pending = self.extend_start(req_id, new_token_ids, new_slot_ids)
        self._eval_with_cache(pending.lazy_token, self._req_caches[req_id])
        return self.extend_finalize(pending)

    def _sync_new_kv_to_pool(
        self,
        cache: list[Any],
        cache_start: int,
        slot_ids: list[int],
    ) -> None:
        """Sync attention KV from contiguous cache to pool at the given slots."""
        if not slot_ids or self._attention_kv_pool is None:
            return
        end = cache_start + len(slot_ids)
        slot_ids_mx = mx.array(slot_ids, dtype=mx.int32)
        # TODO: Standardize ContiguousAttentionKVCache size to avoid transpose
        # Transpose cache (1, n_kv_heads, S, head_dim) to pool (S, n_kv_heads, head_dim)
        k_all = mx.stack(
            [
                cache[layer_idx].keys[0, :, cache_start:end, :].transpose(1, 0, 2)
                for layer_idx in self._cache_layout.attention_layer_indices
            ]
        )
        v_all = mx.stack(
            [
                cache[layer_idx].values[0, :, cache_start:end, :].transpose(1, 0, 2)
                for layer_idx in self._cache_layout.attention_layer_indices
            ]
        )
        self._attention_kv_pool.set_kv_all_layers(slot_ids_mx, k_all, v_all)

    def _sync_decode_kv_to_pool(self, req_id: str) -> None:
        """Sync un-flushed decode KV for *req_id* to the shared pool."""
        if self._attention_kv_pool is None or self._req_to_token_pool is None:
            return
        cache = self._req_caches.get(req_id)
        if cache is None:
            return
        current_offset = self._first_attention_cache(cache).offset
        synced_offset = self._req_synced_offset.get(req_id, 0)
        if current_offset <= synced_offset:
            return
        req_pool_idx = self._req_pool_idx.get(req_id)
        if req_pool_idx is None:
            return
        # Read slot IDs from scheduler's req_to_token_pool
        slot_ids = (
            self._req_to_token_pool.req_to_token[
                req_pool_idx, synced_offset:current_offset
            ]
            .to(dtype=int)
            .tolist()
        )
        self._sync_new_kv_to_pool(cache, synced_offset, slot_ids)
        self._req_synced_offset[req_id] = current_offset

    def flush_all_decode_kv(self) -> None:
        """Sync all active requests' un-flushed decode KV to the pool."""
        if self.disable_radix_cache or self._attention_kv_pool is None:
            return
        for req_id in list(self._req_caches.keys()):
            self._sync_decode_kv_to_pool(req_id)

    def decode_batch(
        self,
        req_ids: list[str],
    ) -> list[int]:
        """Decode one token per request."""
        pending = self.decode_batch_start(req_ids)
        # Evaluate lazy_tokens together with every affected cache buffer so
        # the attention write-then-read ordering is materialised in one
        # kernel submission.
        cache_arrays = self._cache_state_arrays(pending.caches)
        mx.eval(pending.lazy_tokens, *cache_arrays)
        return self.decode_batch_finalize(pending)

    def prefill_start(
        self,
        req_id: str,
        new_token_ids: list[int],
        full_token_ids: list[int],
        prefix_slot_ids: list[int],
        new_slot_ids: list[int],
        req_pool_idx: int,
        req: Any | None = None,
    ) -> MlxPendingPrefill:
        """Queue a prefill forward pass without evaluating.

        Returns an :class:`MlxPendingPrefill` containing the lazy
        next-token ``mx.array`` plus everything needed to commit the
        request in :meth:`prefill_finalize`.  The caller drives the GPU
        by handing ``lazy_token`` (and cache state) to ``mx.async_eval``.
        """
        prefix_len = len(prefix_slot_ids)
        if req is not None:
            req.mamba_last_track_seqlen = None

        if self.disable_radix_cache:
            cache = self._acquire_cache()
            input_ids = mx.array([new_token_ids], dtype=mx.int32)
            model_output = self.model(input_ids, cache=cache)
            logits = self._extract_logits(model_output)
            lazy_token = mx.argmax(logits[:, -1, :], axis=-1)
            return MlxPendingPrefill(
                lazy_token=lazy_token,
                cache=cache,
                req_id=req_id,
                full_token_ids=list(full_token_ids),
                req_pool_idx=req_pool_idx,
                synced_offset=0,
            )

        assert self._attention_kv_pool is not None

        new_token_count = len(new_token_ids)
        track_len = self._select_auxiliary_state_track_len(
            prefix_len=prefix_len,
            new_token_count=new_token_count,
            full_len=len(full_token_ids),
            req=req,
        )

        if prefix_len > 0:
            cache = self._cache_with_pool_backed_attention(prefix_slot_ids, prefix_len)
            pool_backed_attention = True
            restored_auxiliary_state = (
                not self._cache_layout.has_auxiliary_state
                or self._restore_auxiliary_state(req_pool_idx, cache)
            )
            if self._cache_layout.has_auxiliary_state and (
                not restored_auxiliary_state or new_token_count == 0
            ):
                # TODO(MLX): exact full-prefix hits need auxiliary state at
                # prefix_len - 1 to recompute last-token logits. The unified
                # tree stores state at the match boundary today, so use a
                # full-prompt fallback for that edge while still syncing newly
                # allocated attention KV below.
                cache = self._acquire_cache()
                input_ids = mx.array([full_token_ids or new_token_ids], dtype=mx.int32)
                model_output = self.model(input_ids, cache=cache)
                logits = self._extract_logits(model_output)
                lazy_token = mx.argmax(logits[:, -1, :], axis=-1)
                if new_slot_ids:
                    self._sync_new_kv_to_pool(cache, prefix_len, new_slot_ids)
                return MlxPendingPrefill(
                    lazy_token=lazy_token,
                    cache=cache,
                    req_id=req_id,
                    full_token_ids=list(full_token_ids),
                    req_pool_idx=req_pool_idx,
                    synced_offset=prefix_len + len(new_slot_ids),
                )
        else:
            cache = self._acquire_cache()
            pool_backed_attention = False

        if new_token_count > 0:
            track_new_count = track_len - prefix_len if track_len is not None else None
            if track_new_count is not None and 0 < track_new_count < new_token_count:
                input_ids = mx.array([new_token_ids[:track_new_count]], dtype=mx.int32)
                self.model(input_ids, cache=cache)
                self._store_tracked_auxiliary_state(req, cache, track_len)
                if pool_backed_attention:
                    cache = self._materialize_pool_backed_attention(cache)
                    pool_backed_attention = False
                extend_tokens = new_token_ids[track_new_count:]
            else:
                extend_tokens = new_token_ids
        else:
            # Full cache hit - rerun last token to get next-token logits
            extend_tokens = full_token_ids[-1:]
            for c in cache:
                c.offset = max(c.offset - 1, 0)

        input_ids = mx.array([extend_tokens], dtype=mx.int32)
        model_output = self.model(input_ids, cache=cache)
        logits = self._extract_logits(model_output)

        if track_len is not None and track_len == prefix_len + new_token_count:
            self._store_tracked_auxiliary_state(req, cache, track_len)

        last_logits = logits[:, -1, :]
        lazy_token = mx.argmax(last_logits, axis=-1)

        # Convert pool-backed attention KV to contiguous attention KV for decode.
        # This appends a lazy slice-assign onto the forward graph; the
        # arrays get materialised when the caller evaluates lazy_token.
        if pool_backed_attention:
            cache = self._materialize_pool_backed_attention(cache)

        if new_slot_ids:
            self._sync_new_kv_to_pool(cache, prefix_len, new_slot_ids)

        return MlxPendingPrefill(
            lazy_token=lazy_token,
            cache=cache,
            req_id=req_id,
            full_token_ids=list(full_token_ids),
            req_pool_idx=req_pool_idx,
            synced_offset=prefix_len + len(new_slot_ids),
        )

    def prefill_finalize(self, pending: MlxPendingPrefill) -> int:
        """Materialise a pending prefill and commit per-request state.

        Must be called *after* ``pending.lazy_token`` has been handed to
        ``mx.async_eval`` / ``mx.eval``.  ``.item()`` here is blocking on
        that specific lazy scalar.
        """
        next_token = int(pending.lazy_token.item())
        self._req_token_ids[pending.req_id] = list(pending.full_token_ids) + [
            next_token
        ]
        self._req_caches[pending.req_id] = pending.cache
        self._req_pool_idx[pending.req_id] = pending.req_pool_idx
        self._req_synced_offset[pending.req_id] = pending.synced_offset
        self._store_auxiliary_state(pending.req_pool_idx, pending.cache)
        return next_token

    def extend_start(
        self,
        req_id: str,
        new_token_ids: list[int],
        new_slot_ids: list[int],
    ) -> MlxPendingExtend:
        """Queue chunked-prefill continuation without evaluating."""
        assert (
            req_id in self._req_caches
        ), f"extend_start called for unknown request {req_id}"

        cache = self._req_caches[req_id]

        input_ids = mx.array([new_token_ids], dtype=mx.int32)
        model_output = self.model(input_ids, cache=cache)
        logits = self._extract_logits(model_output)
        lazy_token = mx.argmax(logits[:, -1, :], axis=-1)

        if not self.disable_radix_cache and new_slot_ids:
            synced = self._req_synced_offset[req_id]
            self._sync_new_kv_to_pool(cache, synced, new_slot_ids)
            new_synced_offset = synced + len(new_slot_ids)
        else:
            new_synced_offset = self._req_synced_offset.get(req_id, 0)

        return MlxPendingExtend(
            lazy_token=lazy_token,
            req_id=req_id,
            new_token_ids=list(new_token_ids),
            new_synced_offset=new_synced_offset,
        )

    def extend_finalize(self, pending: MlxPendingExtend) -> int:
        """Materialise a pending extend and commit per-request state."""
        next_token = int(pending.lazy_token.item())

        prev_tokens = self._req_token_ids[pending.req_id]
        if prev_tokens:
            prev_tokens.pop()  # remove stale intermediate token
        prev_tokens.extend(pending.new_token_ids)
        prev_tokens.append(next_token)

        self._req_synced_offset[pending.req_id] = pending.new_synced_offset
        self._store_auxiliary_state(
            self._req_pool_idx[pending.req_id],
            self._req_caches[pending.req_id],
        )
        return next_token

    def _extract_model_components(self):
        """Cache embedding, norm, and lm_head for layer-by-layer hybrid forward."""
        root = getattr(self.model, "language_model", self.model)
        text_model = getattr(root, "model", root)
        embed = text_model.embed_tokens
        norm = text_model.norm
        if hasattr(root, "lm_head"):
            lm_head = root.lm_head
        elif hasattr(root, "args") and getattr(root.args, "tie_word_embeddings", False):
            lm_head = text_model.embed_tokens.as_linear
        else:
            lm_head = root.lm_head
        return embed, norm, lm_head

    def _decode_with_hybrid_batching(
        self,
        caches: list[list[Any]],
        batched_input: mx.array,
        req_ids: list[str],
    ) -> mx.array:
        """Layer-by-layer hybrid decode for attention plus auxiliary state.

        Attention layers run with batched hidden states via
        ``BatchedDecodeContext``. Auxiliary layers run batched when their
        native cache implements mlx-lm's merge/extract protocol, otherwise
        they fall back to per-request execution.
        """
        batch_size = len(caches)

        hidden_states = self._model_embed(batched_input)

        ctx = self._build_batched_decode_context(caches, req_ids)
        seq_lens = ctx.seq_lens
        max_offset = max(seq_lens)

        set_context(ctx)
        try:
            for layer_idx in range(self._cache_layout.num_layers):
                layer = self._cache_layout.layers[layer_idx]

                if self._cache_layout.attention_attrs[layer_idx] is not None:
                    shim = AttentionOffsetCache(offset=max_offset)
                    hidden_states = layer(hidden_states, mask=None, cache=shim)
                else:
                    layer_caches = [caches[i][layer_idx] for i in range(batch_size)]
                    hidden_states = self._decode_auxiliary_layer(
                        layer,
                        hidden_states,
                        layer_caches,
                    )
        finally:
            clear_context()

        hidden_states = self._model_norm(hidden_states)
        logits = self._extract_logits(self._model_lm_head(hidden_states))
        return mx.argmax(logits[:, -1, :], axis=-1)

    def _decode_auxiliary_layer(
        self,
        layer: Any,
        hidden_states: mx.array,
        layer_caches: list[Any],
    ) -> mx.array:
        """Decode one auxiliary layer, batching when native cache supports it."""
        if self._can_batch_auxiliary_layer(layer, layer_caches):
            return self._decode_auxiliary_layer_batched(
                layer,
                hidden_states,
                layer_caches,
            )

        results = []
        for i, cache in enumerate(layer_caches):
            results.append(layer(hidden_states[i : i + 1], mask=None, cache=cache))
        return mx.concatenate(results, axis=0)

    @staticmethod
    def _can_batch_auxiliary_layer(layer: Any, layer_caches: list[Any]) -> bool:
        """Return whether an auxiliary layer can run with merged cache state.

        Qwen3.5/Qwen3-Next DeltaNet layers use the mlx-lm DecoderLayer shape
        below with ``ArraysCache``. Its ``merge``/``extract`` helpers can batch
        native state temporarily and split it back to per-request cache objects.
        """
        if not layer_caches:
            return False
        if not (
            getattr(layer, "is_linear", False)
            and hasattr(layer, "input_layernorm")
            and hasattr(layer, "linear_attn")
            and hasattr(layer, "post_attention_layernorm")
            and hasattr(layer, "mlp")
        ):
            return False

        cache_type = type(layer_caches[0])
        if not callable(getattr(cache_type, "merge", None)) or not all(
            isinstance(cache, cache_type) and callable(getattr(cache, "extract", None))
            for cache in layer_caches
        ):
            return False
        return True

    @staticmethod
    def _decode_auxiliary_layer_batched(
        layer: Any,
        hidden_states: mx.array,
        layer_caches: list[Any],
    ) -> mx.array:
        residual = hidden_states
        normed = layer.input_layernorm(hidden_states)

        batched_cache = MlxModelRunner._merge_auxiliary_caches(layer_caches)
        mixed = layer.linear_attn(normed, mask=None, cache=batched_cache)

        extract = getattr(batched_cache, "extract", None)
        if not callable(extract):
            raise RuntimeError(
                f"{type(batched_cache).__name__}.merge() returned a cache "
                "without extract(); cannot split auxiliary decode state"
            )
        for i, cache in enumerate(layer_caches):
            split_cache = extract(i)
            MlxModelRunner._replace_cache_contents(cache, split_cache)

        hidden_states = residual + mixed
        return hidden_states + layer.mlp(layer.post_attention_layernorm(hidden_states))

    @staticmethod
    def _merge_auxiliary_caches(layer_caches: list[Any]) -> Any:
        if MlxModelRunner._can_fast_merge_arrays_cache(layer_caches):
            return MlxModelRunner._fast_merge_arrays_cache(layer_caches)
        return type(layer_caches[0]).merge(layer_caches)

    @staticmethod
    def _can_fast_merge_arrays_cache(layer_caches: list[Any]) -> bool:
        cache_type = type(layer_caches[0])
        if cache_type.__name__ != "ArraysCache":
            return False
        return all(
            type(cache) is cache_type
            and isinstance(getattr(cache, "cache", None), list)
            and getattr(cache, "lengths", None) is None
            and getattr(cache, "left_padding", None) is None
            for cache in layer_caches
        )

    @staticmethod
    def _fast_merge_arrays_cache(layer_caches: list[Any]) -> Any:
        """Merge mlx-lm ArraysCache with concat instead of zero+slice writes."""
        cache_type = type(layer_caches[0])
        merged = cache_type(len(layer_caches[0].cache))
        slots = []
        for slot_idx in range(len(layer_caches[0].cache)):
            values = [cache.cache[slot_idx] for cache in layer_caches]
            first = next((value for value in values if value is not None), None)
            if first is None:
                slots.append(None)
                continue
            slots.append(
                mx.concatenate(
                    [
                        value if value is not None else mx.zeros_like(first)
                        for value in values
                    ],
                    axis=0,
                )
            )
        merged.cache = slots
        return merged

    @staticmethod
    def _replace_cache_contents(cache: Any, new_cache: Any) -> None:
        """Replace cache contents while preserving the original cache object."""
        if type(cache) is type(new_cache) and hasattr(cache, "__dict__"):
            cache.__dict__.clear()
            cache.__dict__.update(new_cache.__dict__)
            return
        if hasattr(cache, "state") and hasattr(new_cache, "state"):
            cache.state = new_cache.state
            return
        raise RuntimeError(
            f"Cannot copy {type(new_cache).__name__} state into "
            f"{type(cache).__name__}"
        )

    def _decode_with_native_cache(
        self,
        caches: list[list[Any]],
        input_ids_by_request: list[mx.array],
    ) -> mx.array:
        lazy_token_list = []
        for input_ids, cache in zip(input_ids_by_request, caches):
            model_output = self.model(input_ids, cache=cache)
            logits = self._extract_logits(model_output)
            lazy_token_list.append(mx.argmax(logits[:, -1, :], axis=-1))
        return (
            lazy_token_list[0]
            if len(lazy_token_list) == 1
            else mx.concatenate(lazy_token_list, axis=0)
        )

    def _decode_with_batched_attention(
        self,
        caches: list[list[Any]],
        batched_input: mx.array,
        req_ids: list[str],
    ) -> mx.array:
        ctx = self._build_batched_decode_context(caches, req_ids)
        seq_lens = ctx.seq_lens
        set_context(ctx)
        try:
            max_offset = max(seq_lens)
            shim_cache = [
                AttentionOffsetCache(offset=max_offset)
                for _ in range(self._cache_layout.num_layers)
            ]
            model_output = self.model(batched_input, cache=shim_cache)
            logits = self._extract_logits(model_output)
            return mx.argmax(logits[:, -1, :], axis=-1)
        finally:
            clear_context()

    def _build_batched_decode_context(
        self,
        caches: list[list[Any]],
        req_ids: list[str],
    ) -> BatchedDecodeContext:
        """Build the shared attention/AOT context for one decode step."""
        return BatchedDecodeContext.from_decode(
            caches=caches,
            req_ids=req_ids,
            aot_kernels=self._aot_kernels,
            kv_pool=self._attention_kv_pool,
            req_pool_idx=self._req_pool_idx,
            req_to_token_pool=self._req_to_token_pool,
            attention_layer_indices=self._cache_layout.attention_layer_indices,
            attention_pool_index_by_layer=(
                self._cache_layout.attention_pool_index_by_layer
            ),
        )

    def decode_batch_start(self, req_ids: list[str]) -> MlxPendingDecode:
        """Queue a decode forward pass without evaluating.

        The caller is responsible for calling ``mx.async_eval`` on the
        returned ``lazy_tokens`` (and optionally per-cache state arrays)
        to kick off GPU work before :meth:`decode_batch_finalize`.
        """
        caches = [self._req_caches[rid] for rid in req_ids]
        last_tokens = [self._req_token_ids[rid][-1] for rid in req_ids]
        batched_input = mx.array(last_tokens, dtype=mx.int32)[:, None]

        if self._cache_layout.has_auxiliary_state:
            lazy_tokens = self._decode_with_hybrid_batching(
                caches, batched_input, list(req_ids)
            )
        else:
            lazy_tokens = self._decode_with_batched_attention(
                caches, batched_input, list(req_ids)
            )

        return MlxPendingDecode(
            lazy_tokens=lazy_tokens,
            req_ids=list(req_ids),
            caches=caches,
        )

    def decode_batch_start_chained(
        self,
        prev: MlxPendingDecode,
    ) -> MlxPendingDecode:
        """Build the next decode step on top of a still-lazy previous decode.

        Feeds ``prev.lazy_tokens`` (an unevaluated ``mx.array`` of shape
        ``(B,)``) as the next step's input ids, reusing
        ``prev.caches`` in-place so that per-layer attention KV writes from
        step N and step N+1 land in the same buffers.  MLX
        tracks the full dependency graph, so once ``mx.async_eval`` is
        called the GPU executes N+1 immediately after N with no gap.

        Caller contract:

        * ``prev`` MUST refer to the same set of requests (same order) as
          the batch the caller intends to run next.  Composition changes
          (finished reqs, new prefills) must break the chain instead.
        * After calling this, finalise ``prev`` BEFORE finalising the
          returned pending: state bookkeeping for step N has to happen
          before step N+1's bookkeeping.
        """
        caches = prev.caches

        # After prev's graph ran, each attention KV cache offset was
        # bumped by one per layer - attention wrapper's `write_token`
        # mutates the Python offset synchronously at graph-build time.
        # So layer-0 offsets reflect the position the NEW token will
        # be written at in step N+1 (and equivalently the RoPE offset).
        batched_input = prev.lazy_tokens[:, None]
        if self._cache_layout.has_auxiliary_state:
            lazy_tokens = self._decode_with_hybrid_batching(
                caches, batched_input, prev.req_ids
            )
        else:
            lazy_tokens = self._decode_with_batched_attention(
                caches, batched_input, prev.req_ids
            )

        return MlxPendingDecode(
            lazy_tokens=lazy_tokens,
            req_ids=prev.req_ids,
            caches=caches,
        )

    def decode_batch_finalize(
        self,
        pending: MlxPendingDecode,
    ) -> list[int]:
        """Materialise a pending decode and update per-request token lists.

        ``pending.lazy_tokens.tolist()`` implicitly blocks until that
        specific lazy array (and its graph ancestors, including the
        per-request cache writes for this step) is evaluated.  The
        caller should have previously handed this pending's lazy_tokens
        to ``mx.async_eval`` (or to a subsequent chained step that will
        be async_eval'd).
        """
        raw = pending.lazy_tokens.tolist()
        if not isinstance(raw, list):
            raw = [raw]
        next_tokens = [int(t) for t in raw]

        for i, rid in enumerate(pending.req_ids):
            self._req_token_ids[rid].append(next_tokens[i])

        self._decode_step_ct += 1
        if self._clear_steps > 0 and self._decode_step_ct % self._clear_steps == 0:
            mx.clear_cache()

        return next_tokens

    def has_request(self, req_id: str) -> bool:
        """Check if a request has active state."""
        return req_id in self._req_caches

    def remove_request(self, req_id: str):
        """Sync remaining decode KV to pool, then release request state."""
        if not self.disable_radix_cache:
            self._sync_decode_kv_to_pool(req_id)

        self._req_token_ids.pop(req_id, None)
        cache = self._req_caches.pop(req_id, None)
        if cache is not None:
            self._release_cache(cache)
        self._req_pool_idx.pop(req_id, None)
        self._req_synced_offset.pop(req_id, None)

    def clear(self):
        """Clear all request states."""
        self._req_token_ids.clear()
        for cache in self._req_caches.values():
            self._release_cache(cache)
        self._req_caches.clear()
        self._req_pool_idx.clear()
        self._req_synced_offset.clear()
        if self._attention_kv_pool is not None:
            self._attention_kv_pool.clear()