From a952e9174f544816325631835b46132c4b17a880 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Sat, 30 May 2026 17:05:02 +0800 Subject: [PATCH] [MLX] Support Qwen3.5 (dense) Model (#25754) Signed-off-by: Xiaodong Ye Co-authored-by: Alex Nails Co-authored-by: Claude Opus 4.6 --- .isort.cfg | 1 + python/pyproject_other.toml | 1 + python/sglang/bench_one_batch.py | 2 +- python/sglang/srt/entrypoints/http_server.py | 5 +- python/sglang/srt/hardware_backend/mlx/aot.py | 15 +- .../hardware_backend/mlx/kv_cache/__init__.py | 44 +- .../mlx/kv_cache/attention_contract.py | 66 + ...tiguous_cache.py => attention_kv_cache.py} | 28 +- .../{kv_pool.py => attention_kv_pool.py} | 18 +- .../mlx/kv_cache/attention_wrapper.py | 96 +- .../mlx/kv_cache/auxiliary_state.py | 390 +++++ .../hardware_backend/mlx/kv_cache/layout.py | 93 ++ .../mlx/kv_cache/model_patching.py | 37 +- .../srt/hardware_backend/mlx/model_runner.py | 823 ++++++++-- .../hardware_backend/mlx/model_runner_stub.py | 41 +- .../hardware_backend/mlx/scheduler_mixin.py | 50 +- .../srt/hardware_backend/mlx/tp_worker.py | 38 +- .../batch_result_processor.py | 5 + .../sglang/srt/mem_cache/cache_init_params.py | 6 + python/sglang/srt/mem_cache/registry.py | 11 +- .../srt/mem_cache/unified_radix_cache.py | 8 +- .../mlx/test_attention_patching.py | 1407 +++++++++++++++++ .../test_unified_radix_cache_unittest.py | 49 + 23 files changed, 2943 insertions(+), 291 deletions(-) create mode 100644 python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py rename python/sglang/srt/hardware_backend/mlx/kv_cache/{contiguous_cache.py => attention_kv_cache.py} (88%) rename python/sglang/srt/hardware_backend/mlx/kv_cache/{kv_pool.py => attention_kv_pool.py} (78%) create mode 100644 python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py create mode 100644 python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py create mode 100644 test/registered/unit/hardware_backend/mlx/test_attention_patching.py diff --git a/.isort.cfg b/.isort.cfg index 835095154..63e968ea6 100644 --- a/.isort.cfg +++ b/.isort.cfg @@ -1,3 +1,4 @@ [settings] profile=black known_first_party=sglang +known_third_party=mlx,mlx_lm diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index b30d80f62..631ecfa31 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -19,6 +19,7 @@ dependencies = ["aiohttp", "requests", "tqdm", "numpy", "IPython", "setproctitle runtime_common = [ "IPython", "aiohttp", + "apache-tvm-ffi", "anthropic>=0.20.0", "blobfile==3.0.0", "av", diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index a8b2e59e1..1022ea3d6 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -536,7 +536,7 @@ class _MlxBenchRunner: if server_args.max_total_tokens is not None: init_kwargs["pool_size"] = server_args.max_total_tokens self.mlx_runner = MlxModelRunner(**init_kwargs) - self.mlx_runner.init_kv_pool(req_to_token_pool=None) + self.mlx_runner.init_cache_pools(req_to_token_pool=None) self.fake_torch_runner = model_runner def clear(self): diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 9c76ae2cc..9dc6c0cde 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -164,6 +164,7 @@ from sglang.srt.utils import ( add_prometheus_track_response_middleware, delete_directory, get_bool_env_var, + is_mps, kill_process_tree, set_uvicorn_logging_configs, ) @@ -1903,8 +1904,8 @@ def _execute_server_warmup(server_args: ServerArgs): model_info = res.json() - # Construct a warmup request - is_vlm = bool(model_info.get("has_image_understanding", False)) + # Construct a warmup request (MLX: text warmup for VLM-advertising models; TODO: enable image warmup). + is_vlm = bool(model_info.get("has_image_understanding", False)) and not is_mps() if model_info["is_generation"]: if is_vlm and not server_args.skip_tokenizer_init: request_name = "/v1/chat/completions" diff --git a/python/sglang/srt/hardware_backend/mlx/aot.py b/python/sglang/srt/hardware_backend/mlx/aot.py index 880514a77..d9be297d2 100644 --- a/python/sglang/srt/hardware_backend/mlx/aot.py +++ b/python/sglang/srt/hardware_backend/mlx/aot.py @@ -13,8 +13,8 @@ from sglang.srt.environ import envs logger = logging.getLogger(__name__) if TYPE_CHECKING: - from sglang.srt.hardware_backend.mlx.kv_cache.contiguous_cache import ( - ContiguousKVCache, + from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import ( + ContiguousAttentionKVCache, ) @@ -123,6 +123,10 @@ class MlxAOTKernelRegistry: 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): @@ -136,10 +140,13 @@ def _build_rope_kernel(inputs: MlxAOTKernelBuildInputs) -> MlxAOTRoPEKernel: 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(sample_attn.n_heads), + "num_qo_heads": int(num_qo_heads), "num_kv_heads": int(inputs.n_kv_heads), } try: @@ -196,7 +203,7 @@ class MlxAOTKernelContext: req_ids: list[str], req_pool_idx: dict[str, int], req_to_token_pool: Any | None, - layer_caches: list[list[ContiguousKVCache]], + 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: diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py index 208e64452..89c83dc94 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/__init__.py @@ -1,5 +1,20 @@ -"""KV cache components for the MLX backend.""" +"""Cache components for the MLX backend.""" +from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import ( + get_head_dim, + get_num_heads, + get_num_kv_heads, + is_attention_module, + uses_sliding_window_attention, +) +from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import ( + AttentionOffsetCache, + ContiguousAttentionKVCache, + PoolBackedAttentionKVCache, +) +from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import ( + MlxAttentionKVPool, +) from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import ( BatchedDecodeContext, MLXAttentionWrapper, @@ -7,12 +22,12 @@ from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import ( get_context, set_context, ) -from sglang.srt.hardware_backend.mlx.kv_cache.contiguous_cache import ( - ContiguousKVCache, - OffsetCache, - PoolBackedCache, +from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import ( + MlxAuxiliaryStateComponent, + MlxAuxiliaryStatePool, + MlxAuxiliaryStateReqToTokenPool, ) -from sglang.srt.hardware_backend.mlx.kv_cache.kv_pool import MlxKVPool +from sglang.srt.hardware_backend.mlx.kv_cache.layout import MlxModelCacheLayout from sglang.srt.hardware_backend.mlx.kv_cache.model_patching import ( find_attention_layers, get_num_layers, @@ -22,14 +37,23 @@ from sglang.srt.hardware_backend.mlx.kv_cache.model_patching import ( __all__ = [ "BatchedDecodeContext", "clear_context", - "ContiguousKVCache", + "AttentionOffsetCache", + "ContiguousAttentionKVCache", "find_attention_layers", + "get_head_dim", "get_context", "get_num_layers", + "get_num_heads", + "get_num_kv_heads", + "is_attention_module", "MLXAttentionWrapper", - "MlxKVPool", - "OffsetCache", + "MlxAttentionKVPool", + "MlxAuxiliaryStateComponent", + "MlxAuxiliaryStatePool", + "MlxAuxiliaryStateReqToTokenPool", + "MlxModelCacheLayout", "patch_model_attention", - "PoolBackedCache", + "PoolBackedAttentionKVCache", "set_context", + "uses_sliding_window_attention", ] diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py new file mode 100644 index 000000000..80f989f1c --- /dev/null +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_contract.py @@ -0,0 +1,66 @@ +"""Attention helpers based on duck typing for the MLX backend.""" + +from __future__ import annotations + +from typing import Any, Iterable + +# ``rope`` and ``scale`` are required by MLXAttentionWrapper. Keeping them in +# the contract also prevents recurrent mixers such as DeltaNet from being +# mistaken for softmax attention just because they expose projection layers. +ATTENTION_API_ATTRS = ("q_proj", "k_proj", "v_proj", "o_proj", "rope", "scale") +NUM_HEAD_ATTRS = ("n_heads", "num_heads", "num_attention_heads") +NUM_KV_HEAD_ATTRS = ("n_kv_heads", "num_k_heads", "num_kv_heads", "num_key_value_heads") +SLIDING_ATTENTION_ATTRS = ( + "is_sliding", + "use_sliding", + "is_sliding_window", + "use_sliding_window", + "is_swa", +) + + +def first_present_attr(module: Any, names: Iterable[str]) -> Any | None: + """Return the first present attribute value without treating 0 as absent.""" + for name in names: + if hasattr(module, name): + return getattr(module, name) + return None + + +def get_num_heads(module: Any) -> int | None: + return first_present_attr(module, NUM_HEAD_ATTRS) + + +def get_num_kv_heads(module: Any) -> int | None: + return first_present_attr(module, NUM_KV_HEAD_ATTRS) + + +def get_head_dim(module: Any) -> int | None: + head_dim = first_present_attr(module, ("head_dim",)) + if head_dim is not None: + return head_dim + + n_kv_heads = get_num_kv_heads(module) + if n_kv_heads is None: + return None + if hasattr(module, "hidden_size") and hasattr(module, "num_k_heads"): + return module.hidden_size // module.num_k_heads + if hasattr(module, "k_proj") and hasattr(module.k_proj, "weight"): + return module.k_proj.weight.shape[0] // n_kv_heads + return None + + +def is_attention_module(module: Any) -> bool: + return ( + all(hasattr(module, attr) for attr in ATTENTION_API_ATTRS) + and get_num_heads(module) is not None + and get_num_kv_heads(module) is not None + ) + + +def uses_sliding_window_attention(*modules: Any) -> bool: + return any( + bool(getattr(module, attr, False)) + for module in modules + for attr in SLIDING_ATTENTION_ATTRS + ) diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/contiguous_cache.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py similarity index 88% rename from python/sglang/srt/hardware_backend/mlx/kv_cache/contiguous_cache.py rename to python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py index 254d2be08..ab0d28ea1 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/contiguous_cache.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_cache.py @@ -1,4 +1,4 @@ -"""ContiguousKVCache, PoolBackedCache and OffsetCache for MLX backend.""" +"""Attention KV cache adapters for the MLX backend.""" from __future__ import annotations @@ -7,10 +7,12 @@ from typing import TYPE_CHECKING import mlx.core as mx if TYPE_CHECKING: - from sglang.srt.hardware_backend.mlx.kv_cache.kv_pool import MlxKVPool + from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_pool import ( + MlxAttentionKVPool, + ) -class OffsetCache: +class AttentionOffsetCache: """Data-free shim satisfying mlx-lm's cache protocol. Provides ``make_mask`` and ``state`` without storing actual K/V. @@ -27,14 +29,14 @@ class OffsetCache: return None if N == 1 else "causal" def update_and_fetch(self, keys, values): - raise RuntimeError("OffsetCache should not store data") + raise RuntimeError("AttentionOffsetCache should not store data") _DEFAULT_MAX_SEQ_LEN = 4096 -class ContiguousKVCache: - """Pre-allocated KV buffer for one request × one layer. +class ContiguousAttentionKVCache: + """Pre-allocated attention KV buffer for one request and one layer. Shape ``(1, n_kv_heads, max_seq_len, head_dim)``. Slice assignment instead of ``mx.concatenate``. Lazy-allocated on first write. @@ -119,12 +121,12 @@ class ContiguousKVCache: return self.keys[:, :, : self.offset, :], self.values[:, :, : self.offset, :] -class PoolBackedCache: - """Lazily gathers cached KV from the shared pool during forward pass. +class PoolBackedAttentionKVCache: + """Lazily gathers cached attention KV from the shared pool during forward. Each ``update_and_fetch`` gathers this layer's prefix from the pool on demand, keeping operations in the lazy compute graph. Convert to - ``ContiguousKVCache`` via ``to_contiguous`` after the forward pass. + ``ContiguousAttentionKVCache`` via ``to_contiguous`` after the forward pass. """ __slots__ = ( @@ -140,7 +142,7 @@ class PoolBackedCache: def __init__( self, - pool: MlxKVPool, + pool: MlxAttentionKVPool, layer_idx: int, slots: mx.array, prefix_len: int, @@ -197,9 +199,9 @@ class PoolBackedCache: self._new_values = values return k_all, v_all - def to_contiguous(self, max_seq_len: int = 4096) -> ContiguousKVCache: - """Convert to ContiguousKVCache reusing forward-pass arrays.""" - cache = ContiguousKVCache(max_seq_len=max_seq_len) + def to_contiguous(self, max_seq_len: int = 4096) -> ContiguousAttentionKVCache: + """Convert to contiguous attention KV reusing forward-pass arrays.""" + cache = ContiguousAttentionKVCache(max_seq_len=max_seq_len) if self._full_keys is not None: cache.update_and_fetch(self._full_keys, self._full_values) return cache diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/kv_pool.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_pool.py similarity index 78% rename from python/sglang/srt/hardware_backend/mlx/kv_cache/kv_pool.py rename to python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_pool.py index 636c20d98..8761438dc 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/kv_pool.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_kv_pool.py @@ -1,4 +1,10 @@ -"""Flat KV pool with per-layer buffers of shape (pool_size, n_kv_heads, head_dim). +"""Flat attention KV pool for the MLX backend. + +Each layer buffer has shape ``(pool_size, n_kv_heads, head_dim)``. +This v1 pool is intentionally uniform: every wrapped softmax-attention +layer must share the same KV shape and full-context KV semantics. +Heterogeneous KV shapes and sliding-window KV need per-layer/window-aware +pools before they can use MLX radix reuse. Slot 0 is reserved as padding (1-based indexing). """ @@ -10,8 +16,8 @@ import mlx.core as mx logger = logging.getLogger(__name__) -class MlxKVPool: - """Pre-allocated KV pool indexed by integer slot IDs.""" +class MlxAttentionKVPool: + """Pre-allocated attention KV pool indexed by integer slot IDs.""" def __init__( self, @@ -27,7 +33,7 @@ class MlxKVPool: self.head_dim = head_dim self.dtype = dtype - # Per-layer buffers: (pool_size, n_kv_heads, head_dim) + # Per-attention-layer buffers: (pool_size, n_kv_heads, head_dim) self.k_buffer: list[mx.array] = [ mx.zeros((pool_size, n_kv_heads, head_dim), dtype=dtype) for _ in range(num_layers) @@ -41,8 +47,8 @@ class MlxKVPool: 1024 * 1024 ) logger.info( - f"MlxKVPool: {pool_size} slots × {num_layers} layers " - f"× {n_kv_heads} heads × {head_dim} dim, " + f"MlxAttentionKVPool: {pool_size} slots x {num_layers} layers " + f"x {n_kv_heads} heads x {head_dim} dim, " f"dtype={dtype}, ~{mem_mb:.1f} MB" ) diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py index dc77bdc53..560b5c820 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/attention_wrapper.py @@ -14,7 +14,14 @@ from sglang.srt.hardware_backend.mlx.aot import ( MlxAOTKernelSet, MlxAOTRoPEContext, ) -from sglang.srt.hardware_backend.mlx.kv_cache.contiguous_cache import ContiguousKVCache +from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import ( + get_head_dim, + get_num_heads, + get_num_kv_heads, +) +from sglang.srt.hardware_backend.mlx.kv_cache.attention_kv_cache import ( + ContiguousAttentionKVCache, +) _thread_local = threading.local() @@ -26,8 +33,9 @@ class BatchedDecodeContext: batch_size: int seq_lens: list[int] # per-request token count before the new token - # layer_caches[layer_idx][req_idx] = ContiguousKVCache - layer_caches: list[list[ContiguousKVCache]] + # attention_layer_caches[attention_pool_idx][req_idx] = ContiguousAttentionKVCache + attention_layer_caches: list[list[ContiguousAttentionKVCache]] + attention_pool_index_by_layer: dict[int, int] = field(default_factory=dict) # Optional AOT kernel state. Keep kernel-specific fields out of the regular # MLX decode path so future AOT kernels can be added without growing this @@ -51,36 +59,46 @@ class BatchedDecodeContext: self.needs_padding = min(seq_lens) < max_seq_len self.pad_sizes = [max_seq_len - s for s in seq_lens] self.positions = mx.arange(self.max_len) if self.needs_padding else None + if not self.attention_pool_index_by_layer: + self.attention_pool_index_by_layer = { + idx: idx for idx in range(len(self.attention_layer_caches)) + } @classmethod def from_decode( cls, *, - caches: list[list[ContiguousKVCache]], - num_layers: int, + caches: list[list[Any]], req_ids: list[str], aot_kernels: MlxAOTKernelSet, kv_pool: Any | None, req_pool_idx: dict[str, int], req_to_token_pool: Any | None, + attention_layer_indices: list[int] | None = None, + attention_pool_index_by_layer: dict[int, int] | None = None, ) -> "BatchedDecodeContext": batch_size = len(req_ids) - seq_lens = [caches[i][0].offset for i in range(batch_size)] - layer_caches = [ + if attention_layer_indices is None: + attention_layer_indices = list(range(len(caches[0]))) + seq_lens = [ + caches[i][attention_layer_indices[0]].offset for i in range(batch_size) + ] + attention_layer_caches = [ [caches[i][layer_idx] for i in range(batch_size)] - for layer_idx in range(num_layers) + for layer_idx in attention_layer_indices ] return cls( batch_size=batch_size, seq_lens=seq_lens, - layer_caches=layer_caches, + attention_layer_caches=attention_layer_caches, + attention_pool_index_by_layer=attention_pool_index_by_layer or {}, aot=MlxAOTKernelContext.from_decode( aot_kernels=aot_kernels, kv_pool=kv_pool, req_ids=req_ids, req_pool_idx=req_pool_idx, req_to_token_pool=req_to_token_pool, - layer_caches=layer_caches, + layer_caches=attention_layer_caches, ), ) @@ -119,15 +137,38 @@ class MLXAttentionWrapper(nn.Module): inner = self._inner layer_idx = self._layer_idx B = ctx.batch_size + n_heads = get_num_heads(inner) + n_kv_heads = get_num_kv_heads(inner) + if n_heads is None or n_kv_heads is None: + raise RuntimeError( + f"Cannot determine attention head counts for {type(inner).__name__}" + ) - queries = inner.q_proj(x) + q_proj_output = inner.q_proj(x) keys = inner.k_proj(x) values = inner.v_proj(x) - head_dim = queries.shape[-1] // inner.n_heads - queries = queries.reshape(B, 1, inner.n_heads, head_dim) - keys = keys.reshape(B, 1, inner.n_kv_heads, head_dim) - values = values.reshape(B, 1, inner.n_kv_heads, head_dim) + head_dim = get_head_dim(inner) + if head_dim is None: + head_dim = keys.shape[-1] // n_kv_heads + + q_width = n_heads * head_dim + gate = None + if q_proj_output.shape[-1] == q_width: + queries = q_proj_output.reshape(B, 1, n_heads, head_dim) + elif q_proj_output.shape[-1] == 2 * q_width: + queries, gate = mx.split( + q_proj_output.reshape(B, 1, n_heads, 2 * head_dim), 2, axis=-1 + ) + gate = gate.reshape(B, 1, q_width) + else: + raise RuntimeError( + f"Unexpected q_proj output shape {q_proj_output.shape} for " + f"{type(inner).__name__}" + ) + + keys = keys.reshape(B, 1, n_kv_heads, head_dim) + values = values.reshape(B, 1, n_kv_heads, head_dim) if hasattr(inner, "q_norm"): queries = inner.q_norm(queries) @@ -140,6 +181,7 @@ class MLXAttentionWrapper(nn.Module): # Vectorized RoPE with per-batch offsets (cached on the context). offsets = ctx.offsets + attention_pool_idx = ctx.attention_pool_index_by_layer[layer_idx] if ctx.aot.rope is not None: # AOT path: real .metallib RoPE + fused KV pool scatter. @@ -148,7 +190,7 @@ class MLXAttentionWrapper(nn.Module): keys, values, offsets, - layer_idx, + attention_pool_idx, ctx.aot.rope, ) else: @@ -157,7 +199,7 @@ class MLXAttentionWrapper(nn.Module): queries = inner.rope(queries, offset=offsets) keys = inner.rope(keys, offset=offsets) - layer_caches = ctx.layer_caches[layer_idx] + layer_caches = ctx.attention_layer_caches[attention_pool_idx] pad_sizes = ctx.pad_sizes # TODO: replace per-request loop with native batched/ragged @@ -173,12 +215,8 @@ class MLXAttentionWrapper(nn.Module): pad = pad_sizes[i] if pad > 0: - k_pad = mx.zeros( - (1, inner.n_kv_heads, pad, head_dim), dtype=k_all.dtype - ) - v_pad = mx.zeros( - (1, inner.n_kv_heads, pad, head_dim), dtype=v_all.dtype - ) + k_pad = mx.zeros((1, n_kv_heads, pad, head_dim), dtype=k_all.dtype) + v_pad = mx.zeros((1, n_kv_heads, pad, head_dim), dtype=v_all.dtype) k_all = mx.concatenate([k_all, k_pad], axis=2) v_all = mx.concatenate([v_all, v_pad], axis=2) @@ -202,6 +240,8 @@ class MLXAttentionWrapper(nn.Module): ) output = output.transpose(0, 2, 1, 3).reshape(B, 1, -1) + if gate is not None: + output = output * mx.sigmoid(gate) return inner.o_proj(output) @staticmethod @@ -210,7 +250,7 @@ class MLXAttentionWrapper(nn.Module): keys: mx.array, values: mx.array, positions: mx.array, - layer_idx: int, + attention_pool_idx: int, rope_ctx: MlxAOTRoPEContext, ) -> tuple[mx.array, mx.array]: """AOT path: rotate Q/K and scatter K/V into the shared pool. @@ -234,8 +274,8 @@ class MLXAttentionWrapper(nn.Module): else: slots = rope_ctx.new_token_slots.astype(mx.int32) - k_pool = rope_ctx.kv_pool.k_buffer[layer_idx] - v_pool = rope_ctx.kv_pool.v_buffer[layer_idx] + k_pool = rope_ctx.kv_pool.k_buffer[attention_pool_idx] + v_pool = rope_ctx.kv_pool.v_buffer[attention_pool_idx] q_rot, k_rot, k_pool_new, v_pool_new = rope_ctx.kernel.rope_pool_fused( q_flat, @@ -251,8 +291,8 @@ class MLXAttentionWrapper(nn.Module): rope_base=rope_ctx.kernel.base, ) # Rebind pool buffers (zero-copy donation result). - rope_ctx.kv_pool.k_buffer[layer_idx] = k_pool_new - rope_ctx.kv_pool.v_buffer[layer_idx] = v_pool_new + rope_ctx.kv_pool.k_buffer[attention_pool_idx] = k_pool_new + rope_ctx.kv_pool.v_buffer[attention_pool_idx] = v_pool_new # (B, n_heads, head_dim) -> (B, n_heads, 1, head_dim) for SDPA path return q_rot[:, :, None, :], k_rot[:, :, None, :] diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py new file mode 100644 index 000000000..be5912193 --- /dev/null +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py @@ -0,0 +1,390 @@ +"""MLX auxiliary-state snapshots for unified radix cache. + +Hybrid MLX models may include non-softmax-attention layers whose native +``mlx-lm`` cache state cannot be reconstructed from the attention KV pool. +The global scheduler exposes that state through its existing MAMBA component +contract, so this MLX adapter keeps those scheduler-facing field names while +storing model-agnostic native cache snapshots. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Iterable, Optional + +import mlx.core as mx +import torch + +from sglang.srt.mem_cache.base_prefix_cache import EvictParams, InsertResult +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.mem_cache.unified_cache_components.mamba_component import ( + MambaComponent, +) +from sglang.srt.mem_cache.unified_cache_components.tree_component import TreeComponent + +_CACHE_ATTRS = ("offset", "lengths", "left_padding") +_MISSING = object() + + +def _clone_tree(value: Any) -> Any: + if isinstance(value, mx.array): + return mx.array(value) + if isinstance(value, list): + return [_clone_tree(item) for item in value] + if isinstance(value, tuple): + return tuple(_clone_tree(item) for item in value) + if isinstance(value, dict): + return {key: _clone_tree(item) for key, item in value.items()} + return value + + +def _arrays_in_tree(value: Any) -> list[mx.array]: + arrays: list[mx.array] = [] + + def collect(item: Any) -> None: + if isinstance(item, mx.array): + arrays.append(item) + elif isinstance(item, (list, tuple)): + for child in item: + collect(child) + elif isinstance(item, dict): + for child in item.values(): + collect(child) + + collect(value) + return arrays + + +@dataclass +class _CacheSnapshot: + state: Any + meta_state: Any = _MISSING + attrs: dict[str, Any] | None = None + + +def _snapshot_cache(cache: Any) -> _CacheSnapshot: + state = _clone_tree(getattr(cache, "state", ())) + meta_state = ( + _clone_tree(cache.meta_state) if hasattr(cache, "meta_state") else _MISSING + ) + attrs = { + name: _clone_tree(getattr(cache, name)) + for name in _CACHE_ATTRS + if hasattr(cache, name) + } + arrays = _arrays_in_tree((state, meta_state, attrs)) + if arrays: + mx.eval(*arrays) + return _CacheSnapshot(state=state, meta_state=meta_state, attrs=attrs) + + +def _restore_cache(cache: Any, snapshot: _CacheSnapshot) -> None: + cache.state = _clone_tree(snapshot.state) + if snapshot.meta_state is not _MISSING and hasattr(cache, "meta_state"): + cache.meta_state = _clone_tree(snapshot.meta_state) + for name, value in (snapshot.attrs or {}).items(): + setattr(cache, name, _clone_tree(value)) + + +class MlxAuxiliaryStatePool: + """Index-addressable snapshots of native MLX auxiliary cache state.""" + + def __init__(self, size: int, device: str): + self.size = size + self.device = device + self.mamba_cache = None + self.mem_usage = 0 + self._snapshots: dict[int, dict[int, _CacheSnapshot]] = {} + self.clear() + + def _tensor(self, indices: Any) -> torch.Tensor: + return torch.as_tensor(indices, dtype=torch.int64, device=self.device).view(-1) + + def _index(self, index: Any) -> int: + flat = self._tensor(index) + assert flat.numel() == 1 + return int(flat.item()) + + def available_size(self) -> int: + return int(self.free_slots.numel()) + + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + if need_size > self.available_size(): + return None + slots = self.free_slots[:need_size].clone() + self.free_slots = self.free_slots[need_size:] + for slot in slots.tolist(): + self._snapshots.pop(int(slot), None) + return slots + + def free(self, indices: Any) -> None: + if indices is None: + return + indices = self._tensor(indices) + if indices.numel() == 0: + return + for slot in indices.tolist(): + self._snapshots.pop(int(slot), None) + self.free_slots = torch.cat([self.free_slots, indices]) + + def clear(self) -> None: + self.free_slots = torch.arange( + 1, self.size + 1, dtype=torch.int64, device=self.device + ) + self._snapshots.clear() + + def copy_from(self, src: Any, dst: Any) -> None: + src_indices = self._tensor(src) + dst_indices = self._tensor(dst) + assert src_indices.numel() == dst_indices.numel() + for src_idx, dst_idx in zip(src_indices.tolist(), dst_indices.tolist()): + snapshot = self._snapshots.get(int(src_idx)) + if snapshot is None: + self._snapshots.pop(int(dst_idx), None) + else: + self._snapshots[int(dst_idx)] = { + layer_idx: _CacheSnapshot( + state=_clone_tree(cache_snapshot.state), + meta_state=_clone_tree(cache_snapshot.meta_state), + attrs=_clone_tree(cache_snapshot.attrs), + ) + for layer_idx, cache_snapshot in snapshot.items() + } + + def fork_from(self, src: Any) -> Optional[torch.Tensor]: + src_indices = self._tensor(src) + dst = self.alloc(src_indices.numel()) + if dst is None: + return None + self.copy_from(src_indices, dst) + return dst + + def store_cache( + self, + index: Any, + cache: list[Any], + layer_indices: Iterable[int], + ) -> None: + self._snapshots[self._index(index)] = { + layer_idx: _snapshot_cache(cache[layer_idx]) for layer_idx in layer_indices + } + + def restore_cache( + self, + index: Any, + cache: list[Any], + layer_indices: Iterable[int] | None = None, + ) -> bool: + snapshot = self._snapshots.get(self._index(index)) + if snapshot is None: + return False + selected_layers = set(layer_indices) if layer_indices is not None else None + for layer_idx, cache_snapshot in snapshot.items(): + if selected_layers is not None and layer_idx not in selected_layers: + continue + _restore_cache(cache[layer_idx], cache_snapshot) + return True + + def has_snapshot(self, index: Any) -> bool: + return self._index(index) in self._snapshots + + +class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool): + """Req-to-token pool with MLX auxiliary-state slot bookkeeping.""" + + def __init__( + self, + *, + size: int, + max_context_len: int, + device: str, + enable_memory_saver: bool, + auxiliary_state_size: int, + ): + super().__init__( + size=size, + max_context_len=max_context_len, + device=device, + enable_memory_saver=enable_memory_saver, + ) + self.mamba_pool = MlxAuxiliaryStatePool( + size=auxiliary_state_size, + device=device, + ) + # The unified radix base MAMBA component still reads ``mamba_pool``. + # Keep the MLX-owned name beside it so local code can avoid model- + # specific terminology. + self.auxiliary_state_pool = self.mamba_pool + self.enable_mamba_extra_buffer = False + self.req_index_to_auxiliary_state_index_mapping = torch.zeros( + self._alloc_size, dtype=torch.int32, device=device + ) + + def alloc(self, reqs): + select_index = super().alloc(reqs) + if select_index is None: + return None + + auxiliary_state_indices = [] + for req in reqs: + if getattr(req, "mamba_pool_idx", None) is not None: + mid = req.mamba_pool_idx + else: + allocated = self.auxiliary_state_pool.alloc(1) + assert allocated is not None, "Not enough MLX auxiliary state slots" + mid = allocated[0] + req.mamba_pool_idx = mid + auxiliary_state_indices.append(mid.to(dtype=torch.int32)) + self.req_index_to_auxiliary_state_index_mapping[select_index] = torch.stack( + auxiliary_state_indices + ) + return select_index + + def get_auxiliary_state_indices(self, req_indices) -> torch.Tensor: + return self.req_index_to_auxiliary_state_index_mapping[req_indices] + + def get_mamba_indices(self, req_indices) -> torch.Tensor: + return self.get_auxiliary_state_indices(req_indices) + + def get_mamba_ping_pong_other_idx(self, mamba_next_track_idx: int) -> int: + return 0 + + def free_mamba_cache(self, req, mamba_ping_pong_track_buffer_to_keep=None): + if getattr(req, "mamba_pool_idx", None) is not None: + self.auxiliary_state_pool.free(req.mamba_pool_idx.unsqueeze(0)) + req.mamba_pool_idx = None + track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None) + if track_buffer is not None: + if mamba_ping_pong_track_buffer_to_keep is None: + self.auxiliary_state_pool.free(track_buffer) + req.mamba_ping_pong_track_buffer = None + req.mamba_next_track_idx = None + + def free_auxiliary_state_cache(self, req, track_buffer_to_keep=None): + self.free_mamba_cache( + req, + mamba_ping_pong_track_buffer_to_keep=track_buffer_to_keep, + ) + + def free(self, req): + super().free(req) + + def clear(self): + super().clear() + self.auxiliary_state_pool.clear() + self.req_index_to_auxiliary_state_index_mapping.zero_() + + +class MlxAuxiliaryStateComponent(MambaComponent): + """Unified radix component for MLX native auxiliary-state snapshots.""" + + def __init__(self, cache, params): + if params.enable_mamba_extra_buffer: + raise NotImplementedError( + "MLX auxiliary-state radix cache does not support " + "enable_mamba_extra_buffer yet." + ) + pool = getattr(cache.req_to_token_pool, "auxiliary_state_pool", None) + if not isinstance(pool, MlxAuxiliaryStatePool): + raise TypeError( + "MlxAuxiliaryStateComponent requires MlxAuxiliaryStatePool, " + f"got {type(pool)}" + ) + TreeComponent.__init__(self, cache, params) + self.enable_mamba_extra_buffer = False + self._mamba_pool_host = None + + @staticmethod + def _tracked_value(req) -> tuple[object | None, bool]: + track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None) + track_len = getattr(req, "mamba_last_track_seqlen", None) + if track_buffer is not None and track_len is not None: + return track_buffer[0].unsqueeze(-1).clone(), True + if getattr(req, "mamba_pool_idx", None) is None: + return None, False + return req.mamba_pool_idx.unsqueeze(-1).clone(), False + + def prepare_for_caching_req( + self, + req, + insert_params, + token_ids_len: int, + is_finished: bool, + ) -> int | None: + cache_len = getattr(req, "mamba_last_track_seqlen", None) + auxiliary_value, uses_track_slot = self._tracked_value(req) + setattr(insert_params, "mlx_auxiliary_state_uses_track_slot", uses_track_slot) + + if auxiliary_value is None: + return 0 if is_finished else None + + if cache_len is None: + cache_len = token_ids_len + if is_finished: + insert_params.mamba_value = auxiliary_value + else: + source_value = auxiliary_value + forked_value = self.cache.req_to_token_pool.auxiliary_state_pool.fork_from( + source_value + ) + if forked_value is None: + self.cache.evict(EvictParams(num_tokens=0, mamba_num=1)) + forked_value = ( + self.cache.req_to_token_pool.auxiliary_state_pool.fork_from( + source_value + ) + ) + assert forked_value is not None, "Can not alloc MLX auxiliary cache" + insert_params.mamba_value = forked_value + return cache_len + + def cleanup_after_caching_req( + self, + req, + is_finished: bool, + insert_result: InsertResult | None = None, + insert_params=None, + ) -> None: + if not is_finished: + if ( + insert_params is not None + and insert_params.mamba_value is not None + and (insert_result is None or insert_result.mamba_exist) + ): + self.cache.req_to_token_pool.auxiliary_state_pool.free( + insert_params.mamba_value + ) + if bool( + getattr(insert_params, "mlx_auxiliary_state_uses_track_slot", False) + ): + track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None) + if track_buffer is not None: + self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer) + req.mamba_ping_pong_track_buffer = None + req.mamba_next_track_idx = None + req.mamba_last_track_seqlen = None + return + + auxiliary_value_exists = ( + insert_result.mamba_exist if insert_result is not None else True + ) + uses_track_slot = bool( + getattr(insert_params, "mlx_auxiliary_state_uses_track_slot", False) + ) + if uses_track_slot: + keep_track_slot = not auxiliary_value_exists + self.cache.req_to_token_pool.free_auxiliary_state_cache( + req, + track_buffer_to_keep=0 if keep_track_slot else None, + ) + elif auxiliary_value_exists: + self.cache.req_to_token_pool.free_auxiliary_state_cache(req) + else: + # The radix tree now owns the live auxiliary-state slot. + track_buffer = getattr(req, "mamba_ping_pong_track_buffer", None) + if track_buffer is not None: + self.cache.req_to_token_pool.auxiliary_state_pool.free(track_buffer) + req.mamba_ping_pong_track_buffer = None + req.mamba_next_track_idx = None + req.mamba_pool_idx = None + req.mamba_last_track_seqlen = None diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py new file mode 100644 index 000000000..ac54bea44 --- /dev/null +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/layout.py @@ -0,0 +1,93 @@ +"""Model cache layout helpers for the MLX backend.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Sequence + + +@dataclass(frozen=True) +class MlxModelCacheLayout: + """Map model layers to MLX cache storage components. + + Attention layers store softmax-attention KV in the MLX attention KV pool. + Auxiliary layers keep native ``mlx-lm`` cache state and are snapshotted by + the MLX auxiliary-state component. + """ + + layers: tuple[Any, ...] + attention_attrs: tuple[str | None, ...] + attention_layer_indices: tuple[int, ...] + auxiliary_layer_indices: tuple[int, ...] + attention_pool_index_by_layer: dict[int, int] + + @classmethod + def from_attention_discovery( + cls, + layers: Sequence[Any], + attention_attrs: Sequence[str | None], + ) -> "MlxModelCacheLayout": + if len(layers) != len(attention_attrs): + raise ValueError( + "Layer count and attention attribute count differ: " + f"{len(layers)} != {len(attention_attrs)}" + ) + + attention_layer_indices = tuple( + idx for idx, attr in enumerate(attention_attrs) if attr is not None + ) + auxiliary_layer_indices = tuple( + idx for idx, attr in enumerate(attention_attrs) if attr is None + ) + attention_pool_index_by_layer = { + layer_idx: pool_idx + for pool_idx, layer_idx in enumerate(attention_layer_indices) + } + + return cls( + layers=tuple(layers), + attention_attrs=tuple(attention_attrs), + attention_layer_indices=attention_layer_indices, + auxiliary_layer_indices=auxiliary_layer_indices, + attention_pool_index_by_layer=attention_pool_index_by_layer, + ) + + @property + def num_layers(self) -> int: + return len(self.layers) + + @property + def num_attention_layers(self) -> int: + return len(self.attention_layer_indices) + + @property + def has_auxiliary_state(self) -> bool: + return bool(self.auxiliary_layer_indices) + + @property + def first_attention_layer_index(self) -> int: + if not self.attention_layer_indices: + raise RuntimeError("MLX model has no supported attention layers") + return self.attention_layer_indices[0] + + def attention_pool_index(self, layer_idx: int) -> int: + try: + return self.attention_pool_index_by_layer[layer_idx] + except KeyError as exc: + raise KeyError(f"Layer {layer_idx} is not an attention layer") from exc + + def attention_attr(self, layer_idx: int) -> str: + attr = self.attention_attrs[layer_idx] + if attr is None: + raise KeyError(f"Layer {layer_idx} is not an attention layer") + return attr + + def attention_layer_caches( + self, + caches_by_request: list[list[Any]], + ) -> list[list[Any]]: + """Return layer-major attention caches for batched decode.""" + return [ + [request_cache[layer_idx] for request_cache in caches_by_request] + for layer_idx in self.attention_layer_indices + ] diff --git a/python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py b/python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py index 1d1b5065f..ed41317fc 100644 --- a/python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py +++ b/python/sglang/srt/hardware_backend/mlx/kv_cache/model_patching.py @@ -2,25 +2,38 @@ from typing import Any +import mlx.nn as nn + +from sglang.srt.hardware_backend.mlx.kv_cache.attention_contract import ( + is_attention_module, +) from sglang.srt.hardware_backend.mlx.kv_cache.attention_wrapper import ( MLXAttentionWrapper, ) -def find_attention_layers(model: Any) -> tuple[list[Any], str]: - """Find transformer layers and the attention attribute name.""" +def _find_attention_attr(layer: Any) -> str | None: + """Return the direct child name that satisfies the attention contract.""" + if not isinstance(layer, nn.Module): + raise TypeError(f"Expected mlx.nn.Module layer, got {type(layer)}") + for name, module in layer.children().items(): + if isinstance(module, MLXAttentionWrapper) or is_attention_module(module): + return name + return None + + +def find_attention_layers(model: Any) -> tuple[list[Any], list[str | None]]: + """Find transformer layers and per-layer attention attribute names.""" root = getattr(model, "language_model", model) container = getattr(root, "model", root) layer_list = getattr(container, "layers", None) or getattr(root, "layers", []) if layer_list: - sample = layer_list[0] - if hasattr(sample, "self_attn"): - return layer_list, "self_attn" - if hasattr(sample, "attention"): - return layer_list, "attention" - raise ValueError(f"No attention attribute in layer type {type(sample)}") - return layer_list, "self_attn" + attn_attrs = [_find_attention_attr(layer) for layer in layer_list] + if any(attr is not None for attr in attn_attrs): + return layer_list, attn_attrs + raise ValueError(f"No attention attribute in layer type {type(layer_list[0])}") + return layer_list, [] def patch_model_attention(model: Any) -> int: @@ -29,9 +42,11 @@ def patch_model_attention(model: Any) -> int: The wrapper delegates to the inner module when no BatchedDecodeContext is set, so it is always installed and never removed. """ - layer_list, attn_attr = find_attention_layers(model) + layer_list, attn_attrs = find_attention_layers(model) patched = 0 - for idx, layer in enumerate(layer_list): + for idx, (layer, attn_attr) in enumerate(zip(layer_list, attn_attrs)): + if attn_attr is None: + continue attn = getattr(layer, attn_attr) if isinstance(attn, MLXAttentionWrapper): continue diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner.py b/python/sglang/srt/hardware_backend/mlx/model_runner.py index 61f4c086c..5a98730b2 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner.py @@ -2,9 +2,11 @@ Slot allocation and radix-trie prefix matching are handled by the scheduler (``TokenToKVPoolAllocator`` / ``RadixCache``). This runner -reads cached KV from ``MlxKVPool``, runs the forward pass, and writes -new KV back. Each request also keeps a ``ContiguousKVCache`` for -decode-time attention. +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 @@ -17,6 +19,7 @@ state. import logging import time from dataclasses import dataclass +from typing import Any import mlx.core as mx import psutil @@ -29,19 +32,23 @@ from sglang.srt.hardware_backend.mlx.aot import ( MlxAOTKernelSet, ) from sglang.srt.hardware_backend.mlx.kv_cache import ( + AttentionOffsetCache, BatchedDecodeContext, - ContiguousKVCache, + ContiguousAttentionKVCache, + MlxAttentionKVPool, MLXAttentionWrapper, - OffsetCache, - PoolBackedCache, + MlxModelCacheLayout, + PoolBackedAttentionKVCache, clear_context, find_attention_layers, - get_num_layers, + get_head_dim, + get_num_kv_heads, patch_model_attention, set_context, + uses_sliding_window_attention, ) -from sglang.srt.hardware_backend.mlx.kv_cache.kv_pool import MlxKVPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.server_args import get_global_server_args logger = logging.getLogger(__name__) @@ -50,14 +57,14 @@ logger = logging.getLogger(__name__) class MlxPendingPrefill: """Lazy prefill state, finalised after ``mx.eval``/``async_eval``. - ``cache`` is the per-layer list of ``ContiguousKVCache`` that will + ``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 a transient ``PoolBackedCache`` list - already (so its ``state`` arrays are safe to hand to ``async_eval``). + 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 # list[ContiguousKVCache] + cache: list[Any] req_id: str full_token_ids: list[int] req_pool_idx: int @@ -84,7 +91,7 @@ class MlxPendingExtend: class MlxPendingDecode: """Lazy decode state, finalised after ``mx.eval``/``async_eval``. - ``caches`` is a per-request list of per-layer ``ContiguousKVCache`` + ``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 @@ -93,7 +100,7 @@ class MlxPendingDecode: lazy_tokens: mx.array req_ids: list[str] - caches: list # list[list[ContiguousKVCache]] + caches: list[list[Any]] _MLX_QUANTIZATION_PRESETS: dict[str, tuple[int, int]] = { @@ -101,6 +108,7 @@ _MLX_QUANTIZATION_PRESETS: dict[str, tuple[int, int]] = { "mlx_q4": (4, 64), "mlx_q8": (8, 64), } +_MLX_KV_FLOAT_DTYPES = {mx.float16, mx.bfloat16, mx.float32} class MlxModelRunner: @@ -123,9 +131,9 @@ class MlxModelRunner: # Counter used to trigger periodic mx.clear_cache() calls. self._decode_step_ct: int = 0 # On-the-fly quantization preset (e.g. "mlx_q4"). None = no on-load quantization. - # Pre-quantized HF repos (e.g. mlx-community/Qwen3-0.6B-4bit) load correctly - # regardless of this setting — mlx_lm.load() detects the config and instantiates - # QuantizedLinear modules directly. + # 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() @@ -139,14 +147,30 @@ class MlxModelRunner: patch_model_attention(self.model) - self._num_layers = get_num_layers(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[ContiguousKVCache | PoolBackedCache]] = {} + self._req_caches: dict[str, list[Any]] = {} self._req_token_ids: dict[str, list[int]] = {} - self._cache_pool: list[list[ContiguousKVCache]] = [] # reusable caches + self._cache_pool: list[list[Any]] = [] # reusable full-attention caches - self._kv_pool: MlxKVPool | None = None + 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] = {} @@ -161,33 +185,205 @@ class MlxModelRunner: return model_output[0] return model_output - def _acquire_cache(self) -> list[ContiguousKVCache]: + 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 self._cache_pool: + 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 [ - ContiguousKVCache(max_seq_len=self._max_seq_len) - for _ in range(self._num_layers) - ] + return self._new_native_cache() - def _release_cache(self, cache: list[ContiguousKVCache]) -> None: + def _release_cache(self, cache: list[Any]) -> None: """Return a cache list to the pool for reuse.""" - self._cache_pool.append(cache) + if not self._cache_layout.has_auxiliary_state: + self._cache_pool.append(cache) - @staticmethod - def _eval_with_cache( - token_result: mx.array, cache: list[ContiguousKVCache | PoolBackedCache] + 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_global_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: - """Evaluate token result and all cache buffers in one mx.eval call.""" - mx.eval(token_result, *[s for c in cache for s in c.state]) + 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_state_arrays( - pending_caches: list[list[ContiguousKVCache | PoolBackedCache]], - ) -> list[mx.array]: + 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``. @@ -196,7 +392,7 @@ class MlxModelRunner: s for cache_list in pending_caches for cache in cache_list - for s in cache.state + for s in MlxModelRunner._cache_arrays(cache) ] def _load_model(self): @@ -259,38 +455,78 @@ class MlxModelRunner: ) # Force-evaluate weights so mx.get_active_memory() reflects - # actual usage before KV pool sizing. + # 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") - def _get_attn_config(self) -> tuple[int, int, mx.Dtype]: - """Return (n_kv_heads, head_dim, dtype) from the model.""" - layer_list, attn_attr = find_attention_layers(self.model) - if not layer_list: - raise RuntimeError("Cannot determine attention config: no layers found") - sample_attn = getattr(layer_list[0], attn_attr) - if isinstance(sample_attn, MLXAttentionWrapper): - sample_attn = sample_attn._inner - n_kv_heads = sample_attn.n_kv_heads - if hasattr(sample_attn, "head_dim"): - head_dim = sample_attn.head_dim - elif hasattr(sample_attn, "k_proj") and hasattr(sample_attn.k_proj, "weight"): - head_dim = sample_attn.k_proj.weight.shape[0] // n_kv_heads - else: - raise RuntimeError("Cannot determine head_dim from attention module") + 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 stores packed weights as integers, while the KV + # cache stores dequantized projection outputs. + 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._num_layers + 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", @@ -305,7 +541,7 @@ class MlxModelRunner: 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 KV pool: " + 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, " @@ -320,10 +556,13 @@ class MlxModelRunner: def _build_aot_kernels(self) -> MlxAOTKernelSet: """Build model-level set of optional registered AOT kernels.""" - layer_list, attn_attr = find_attention_layers(self.model) - if not layer_list: + if self._cache_layout.num_attention_layers == 0: return MlxAOTKernelSet() - sample_attn = getattr(layer_list[0], attn_attr) + 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, @@ -331,24 +570,25 @@ class MlxModelRunner: head_dim=int(head_dim), ) - def init_kv_pool(self, req_to_token_pool: ReqToTokenPool) -> None: - """Create MlxKVPool (+1 for padding slot 0) and wire scheduler pools.""" + 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._kv_pool = MlxKVPool( + self._attention_kv_pool = MlxAttentionKVPool( pool_size=self._pool_size + 1, - num_layers=self._num_layers, + num_layers=self._cache_layout.num_attention_layers, n_kv_heads=n_kv_heads, head_dim=head_dim, dtype=dtype, ) logger.info( - f"KV pool initialized: pool_size={self._pool_size} " + f"Attention KV pool initialized: pool_size={self._pool_size} " f"(buffer size {self._pool_size + 1} incl. padding slot 0), " - f"{self._num_layers} layers, {n_kv_heads} kv_heads, {head_dim} head_dim" + f"{self._cache_layout.num_attention_layers} attention layers, " + f"{n_kv_heads} kv_heads, {head_dim} head_dim" ) def prefill( @@ -359,6 +599,7 @@ class MlxModelRunner: 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( @@ -368,6 +609,7 @@ class MlxModelRunner: 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) @@ -385,40 +627,39 @@ class MlxModelRunner: def _sync_new_kv_to_pool( self, - cache: list[ContiguousKVCache], + cache: list[Any], cache_start: int, slot_ids: list[int], ) -> None: - """Sync KV from contiguous cache to pool at the given slot IDs.""" - if not slot_ids or self._kv_pool is 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 - num_layers = len(cache) end = cache_start + len(slot_ids) slot_ids_mx = mx.array(slot_ids, dtype=mx.int32) - # TODO: Standardize ContiguousKVCache size to avoid transpose + # 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[i].keys[0, :, cache_start:end, :].transpose(1, 0, 2) - for i in range(num_layers) + 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[i].values[0, :, cache_start:end, :].transpose(1, 0, 2) - for i in range(num_layers) + cache[layer_idx].values[0, :, cache_start:end, :].transpose(1, 0, 2) + for layer_idx in self._cache_layout.attention_layer_indices ] ) - self._kv_pool.set_kv_all_layers(slot_ids_mx, k_all, v_all) + 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._kv_pool is None or self._req_to_token_pool is None: + 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 = cache[0].offset + current_offset = self._first_attention_cache(cache).offset synced_offset = self._req_synced_offset.get(req_id, 0) if current_offset <= synced_offset: return @@ -438,7 +679,7 @@ class MlxModelRunner: 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._kv_pool is None: + 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) @@ -464,6 +705,7 @@ class MlxModelRunner: 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. @@ -472,8 +714,9 @@ class MlxModelRunner: request in :meth:`prefill_finalize`. The caller drives the GPU by handing ``lazy_token`` (and cache state) to ``mx.async_eval``. """ - num_layers = self._num_layers 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() @@ -490,21 +733,62 @@ class MlxModelRunner: synced_offset=0, ) - assert self._kv_pool is not None + 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: - slot_ids_mx = mx.array(prefix_slot_ids, dtype=mx.int32) - cache = [ - PoolBackedCache(self._kv_pool, i, slot_ids_mx, prefix_len) - for i in range(num_layers) - ] + 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: - extend_tokens = new_token_ids + 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:] @@ -515,20 +799,17 @@ class MlxModelRunner: 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 PoolBackedCache to ContiguousKVCache for decode. + # 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 prefix_len > 0: - contiguous_cache = self._acquire_cache() - for layer_idx in range(num_layers): - pbc = cache[layer_idx] - contiguous_cache[layer_idx].update_and_fetch( - pbc._full_keys, pbc._full_values - ) - cache = contiguous_cache + 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) @@ -556,6 +837,7 @@ class MlxModelRunner: 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( @@ -601,8 +883,254 @@ class MlxModelRunner: 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. @@ -610,45 +1138,18 @@ class MlxModelRunner: returned ``lazy_tokens`` (and optionally per-cache state arrays) to kick off GPU work before :meth:`decode_batch_finalize`. """ - batch_size = len(req_ids) - num_layers = self._num_layers - 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 batch_size == 1: - cache = caches[0] - last_token = self._req_token_ids[req_ids[0]][-1] - input_ids = mx.array([[last_token]], dtype=mx.int32) - model_output = self.model(input_ids, cache=cache) - logits = self._extract_logits(model_output) - lazy_tokens = mx.argmax(logits[:, -1, :], axis=-1) - return MlxPendingDecode( - lazy_tokens=lazy_tokens, - req_ids=list(req_ids), - caches=caches, + 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) ) - - ctx = BatchedDecodeContext.from_decode( - caches=caches, - num_layers=num_layers, - req_ids=req_ids, - aot_kernels=self._aot_kernels, - kv_pool=self._kv_pool, - req_pool_idx=self._req_pool_idx, - req_to_token_pool=self._req_to_token_pool, - ) - seq_lens = ctx.seq_lens - set_context(ctx) - try: - max_offset = max(seq_lens) - shim_cache = [OffsetCache(offset=max_offset) for _ in range(num_layers)] - last_tokens = [self._req_token_ids[rid][-1] for rid in req_ids] - batched_input = mx.array(last_tokens, dtype=mx.int32)[:, None] - model_output = self.model(batched_input, cache=shim_cache) - logits = self._extract_logits(model_output) - lazy_tokens = mx.argmax(logits[:, -1, :], axis=-1) - finally: - clear_context() return MlxPendingDecode( lazy_tokens=lazy_tokens, @@ -664,8 +1165,8 @@ class MlxModelRunner: 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 the per-layer ``ContiguousKVCache`` - writes from step N and step N+1 land in the same buffers. MLX + ``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. @@ -678,52 +1179,26 @@ class MlxModelRunner: returned pending: state bookkeeping for step N has to happen before step N+1's bookkeeping. """ - batch_size = len(prev.req_ids) - num_layers = self._num_layers caches = prev.caches - # TODO (changminbark): Need to fix ContiguousKVCache.write_token - # to accommodate dynamic growing like ContiguousKVCache.update_and_fetch. + # TODO (changminbark): Need to fix + # ContiguousAttentionKVCache.write_token to accommodate dynamic growing + # like ContiguousAttentionKVCache.update_and_fetch. - # After prev's graph ran, each ContiguousKVCache.offset was + # 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). - seq_lens = [caches[i][0].offset for i in range(batch_size)] - - if batch_size == 1: - cache = caches[0] - batched_input = prev.lazy_tokens[:, None] - model_output = self.model(batched_input, cache=cache) - logits = self._extract_logits(model_output) - lazy_tokens = mx.argmax(logits[:, -1, :], axis=-1) - return MlxPendingDecode( - lazy_tokens=lazy_tokens, - req_ids=prev.req_ids, - caches=caches, + 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 ) - - ctx = BatchedDecodeContext.from_decode( - caches=caches, - num_layers=num_layers, - req_ids=prev.req_ids, - aot_kernels=self._aot_kernels, - kv_pool=self._kv_pool, - req_pool_idx=self._req_pool_idx, - req_to_token_pool=self._req_to_token_pool, - ) - seq_lens = ctx.seq_lens - set_context(ctx) - try: - max_offset = max(seq_lens) - shim_cache = [OffsetCache(offset=max_offset) for _ in range(num_layers)] - batched_input = prev.lazy_tokens[:, None] - model_output = self.model(batched_input, cache=shim_cache) - logits = self._extract_logits(model_output) - lazy_tokens = mx.argmax(logits[:, -1, :], axis=-1) - finally: - clear_context() return MlxPendingDecode( lazy_tokens=lazy_tokens, @@ -779,9 +1254,9 @@ class MlxModelRunner: """Clear all request states.""" self._req_token_ids.clear() for cache in self._req_caches.values(): - self._cache_pool.append(cache) + self._release_cache(cache) self._req_caches.clear() self._req_pool_idx.clear() self._req_synced_offset.clear() - if self._kv_pool is not None: - self._kv_pool.clear() + if self._attention_kv_pool is not None: + self._attention_kv_pool.clear() diff --git a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py index 3dd8520f2..17f01e93e 100644 --- a/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py +++ b/python/sglang/srt/hardware_backend/mlx/model_runner_stub.py @@ -9,6 +9,9 @@ from typing import Tuple import torch +from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import ( + MlxAuxiliaryStateReqToTokenPool, +) from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool from sglang.srt.model_executor.model_runner import ModelRunner @@ -17,11 +20,11 @@ logger = logging.getLogger(__name__) class _DummyKVCache(KVCache): - """A KV cache that allocates no GPU memory. + """Scheduler-facing KV cache that allocates no GPU memory. Satisfies the KVCache interface so that TokenToKVPoolAllocator can be - constructed, but every buffer access raises — the MLX backend manages - its own KV cache internally. + constructed, but every buffer access raises. The MLX backend manages + attention KV and auxiliary state internally. """ def __init__(self, size: int, dtype: torch.dtype, device: str): @@ -42,16 +45,16 @@ class _DummyKVCache(KVCache): self.custom_mem_pool = None def get_key_buffer(self, layer_id: int) -> torch.Tensor: - raise RuntimeError("_DummyKVCache has no key buffer (MLX manages KV cache)") + raise RuntimeError("_DummyKVCache has no key buffer (MLX manages cache)") def get_value_buffer(self, layer_id: int) -> torch.Tensor: - raise RuntimeError("_DummyKVCache has no value buffer (MLX manages KV cache)") + raise RuntimeError("_DummyKVCache has no value buffer (MLX manages cache)") def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]: - raise RuntimeError("_DummyKVCache has no kv buffer (MLX manages KV cache)") + raise RuntimeError("_DummyKVCache has no kv buffer (MLX manages cache)") def set_kv_buffer(self, layer, loc, cache_k, cache_v) -> None: - raise RuntimeError("_DummyKVCache cannot set kv buffer (MLX manages KV cache)") + raise RuntimeError("_DummyKVCache cannot set kv buffer (MLX manages cache)") def get_kv_size_bytes(self): return 0, 0 @@ -141,12 +144,24 @@ class MlxModelRunnerStub(ModelRunner): self.is_hybrid_swa = False # Create minimal pools - self.req_to_token_pool = ReqToTokenPool( - size=self.max_running_requests, - max_context_len=self.model_config.context_len, - device="cpu", - enable_memory_saver=False, - ) + if self.mambaish_config is not None: + auxiliary_state_size = self.server_args.max_mamba_cache_size + if auxiliary_state_size is None: + auxiliary_state_size = self.max_running_requests * 4 + self.req_to_token_pool = MlxAuxiliaryStateReqToTokenPool( + size=self.max_running_requests, + max_context_len=self.model_config.context_len, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=auxiliary_state_size, + ) + else: + self.req_to_token_pool = ReqToTokenPool( + size=self.max_running_requests, + max_context_len=self.model_config.context_len, + device="cpu", + enable_memory_saver=False, + ) dummy_kv = _DummyKVCache( size=self.max_total_num_tokens, diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index 2f03b9fa1..673aafda2 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -6,11 +6,11 @@ the scheduler runs its CPU-side bookkeeping on the tokens of the older one. The lazy-graph primitives live in ``hardware_backend/mlx/tp_worker.py`` and ``model_runner.py``. -Each request's KV lives ina set of per-request, per-layer ``ContiguousKVCache`` -objects that the ``MLXAttentionWrapper`` mutates in place during the forward pass. -Chained decodes reuse the same cache objects: step N+1's graph reads -step N's lazy writes via MLX's dependency tracking, so the GPU runs -both steps back-to-back with no idle gap. +Each request's attention KV lives in per-request, per-layer +``ContiguousAttentionKVCache`` objects that ``MLXAttentionWrapper`` mutates +in place during the forward pass. Chained decodes reuse the same cache objects: +step N+1's graph reads step N's lazy writes via MLX's dependency tracking, so +the GPU runs both steps back-to-back with no idle gap. """ from __future__ import annotations @@ -62,6 +62,9 @@ class MlxPendingJob: time. Decoupled from the live batch so ``process_batch_result`` can update request state without racing against the next scheduling decision. + schedule_batch: The full scheduler batch. Unlike ``batch_copy``, + this keeps allocator/cache fields needed when a prefill batch + becomes the next running decode batch. reqs: Snapshot of ``batch.reqs`` at launch time. The overlap loop uses this to check ``req.finished()`` on the previous step's request list without holding a reference to the @@ -74,12 +77,27 @@ class MlxPendingJob: decode: Optional["MlxPendingDecode"] mode: str batch_copy: "ScheduleBatch" + schedule_batch: "ScheduleBatch" reqs: List[Req] class SchedulerMlxOverlapMixin: """Mixin that adds MLX overlap scheduling to :class:`Scheduler`.""" + def _finalize_mlx_pending_job(self: "Scheduler", pending: MlxPendingJob): + result = self.tp_worker.finalize_mlx_result( + pending.prefills, + pending.extends, + pending.decode, + pending.mode, + pending.reqs, + ) + if result.next_token_ids is not None: + pending.batch_copy.input_ids = result.next_token_ids + pending.schedule_batch.input_ids = result.next_token_ids + self.last_batch = pending.schedule_batch + self.process_batch_result(pending.batch_copy, result) + @DynamicGradMode() def event_loop_overlap_mlx(self: "Scheduler"): """MLX-specific overlap loop modelled on ``mlx_lm.generate.generate_step``. @@ -123,18 +141,6 @@ class SchedulerMlxOverlapMixin: pending_curr: Optional[MlxPendingJob] = None pending_next: Optional[MlxPendingJob] = None - def _finalize(pending: MlxPendingJob): - result = self.tp_worker.finalize_mlx_result( - pending.prefills, - pending.extends, - pending.decode, - pending.mode, - pending.reqs, - ) - if result.next_token_ids is not None: - pending.batch_copy.input_ids = result.next_token_ids - self.process_batch_result(pending.batch_copy, result) - def _launch_fresh(batch: "ScheduleBatch") -> MlxPendingJob: lazy_tokens, prefills, extends, decode, mode = ( self.tp_worker.async_forward_batch_generation_mlx(batch) @@ -146,6 +152,7 @@ class SchedulerMlxOverlapMixin: decode=decode, mode=mode, batch_copy=batch.copy(), + schedule_batch=batch, reqs=list(batch.reqs), ) @@ -164,6 +171,7 @@ class SchedulerMlxOverlapMixin: decode=decode, mode=mode, batch_copy=prev.batch_copy.copy(), + schedule_batch=prev.schedule_batch, reqs=prev.reqs, ) @@ -191,7 +199,7 @@ class SchedulerMlxOverlapMixin: # 2. Finalize/process on pending_curr's tokens. (GPU is already # executing pending_next at this point.) if pending_curr is not None: - _finalize(pending_curr) + self._finalize_mlx_pending_job(pending_curr) self.result_queue.popleft() pending_curr = None @@ -208,8 +216,8 @@ class SchedulerMlxOverlapMixin: ): pending_curr = pending_next pending_next = None - self.cur_batch = pending_curr.batch_copy - self.last_batch = pending_curr.batch_copy + self.cur_batch = pending_curr.schedule_batch + self.last_batch = pending_curr.schedule_batch if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.invariant_checker.self_check_during_busy() continue @@ -217,7 +225,7 @@ class SchedulerMlxOverlapMixin: # 4. Chain is broken. Finalise pending_next (if any), then # schedule fresh. if pending_next is not None: - _finalize(pending_next) + self._finalize_mlx_pending_job(pending_next) self.result_queue.popleft() pending_next = None next_batch = self.get_next_batch_to_run() diff --git a/python/sglang/srt/hardware_backend/mlx/tp_worker.py b/python/sglang/srt/hardware_backend/mlx/tp_worker.py index c4ecb05aa..3e2f93c2f 100644 --- a/python/sglang/srt/hardware_backend/mlx/tp_worker.py +++ b/python/sglang/srt/hardware_backend/mlx/tp_worker.py @@ -2,7 +2,7 @@ Routes forward passes through the MLX model runner, bypassing PyTorch MPS. A lightweight stub provides scheduler bookkeeping; the actual -KV data lives in MlxKVPool. +attention KV data lives in MlxAttentionKVPool. The worker also exposes an async (lazy-eval) surface used by the MLX overlap scheduler: ``async_forward_batch_generation_mlx`` launches a @@ -87,9 +87,9 @@ class MlxTpModelWorker(TpModelWorker): return None def _ensure_mlx_pool_initialized(self): - """Lazily initialize the MlxKVPool after the stub's pools are ready.""" + """Lazily initialize MLX cache pools after the stub pools are ready.""" if not self._mlx_pool_initialized: - self._mlx_runner.init_kv_pool(self._model_runner.req_to_token_pool) + self._mlx_runner.init_cache_pools(self._model_runner.req_to_token_pool) self._mlx_pool_initialized = True def forward_batch_generation( @@ -124,6 +124,14 @@ class MlxTpModelWorker(TpModelWorker): else: self._mlx_active_rids |= current_rids + def prepare_for_kv_cache_release(self, req) -> None: + """Snapshot MLX auxiliary state at the scheduler's radix insert point.""" + if self._mlx_runner.has_request(req.rid): + self._mlx_runner.store_auxiliary_state_for_request(req.rid) + # Prefer the just-snapshotted live auxiliary state for the final + # insert. Any older tracked slot is released during component cleanup. + req.mamba_last_track_seqlen = None + def _forward_batch_generation_mlx( self, batch: ScheduleBatch ) -> GenerationBatchResult: @@ -144,7 +152,7 @@ class MlxTpModelWorker(TpModelWorker): next_token_ids_list: list[int] = [] if forward_mode.is_extend(): - # Ensure pool is up-to-date before PoolBackedCache reads it + # Ensure pool is up-to-date before pool-backed attention reads it # for prefix-cached prefills. Only runs on extend batches. self._mlx_runner.flush_all_decode_kv() input_ids_cpu = batch.input_ids.cpu().tolist() @@ -185,6 +193,7 @@ class MlxTpModelWorker(TpModelWorker): prefix_slot_ids=prefix_slot_ids, new_slot_ids=req_new_slots, req_pool_idx=req.req_pool_idx, + req=req, ) prefill_rids.append((req.rid, next_token)) @@ -270,7 +279,7 @@ class MlxTpModelWorker(TpModelWorker): if forward_mode.is_extend(): # TODO (changminbark): Implement per-batch flushing using prefix_slot_ids - # Ensure the pool is up-to-date before any PoolBackedCache + # Ensure the pool is up-to-date before pool-backed attention # reads it for prefix-cached prefills. Mirror the sync path. self._mlx_runner.flush_all_decode_kv() return self._async_extend_batch(batch) @@ -330,6 +339,7 @@ class MlxTpModelWorker(TpModelWorker): prefix_slot_ids=prefix_slot_ids, new_slot_ids=req_new_slots, req_pool_idx=req.req_pool_idx, + req=req, ) ) @@ -377,7 +387,23 @@ class MlxTpModelWorker(TpModelWorker): @staticmethod def _cache_state(cache_list) -> list[mx.array]: """Flatten a per-layer cache list to its ``state`` arrays.""" - return [s for c in cache_list for s in c.state] + arrays: list[mx.array] = [] + + def collect(value): + 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) + + for cache in cache_list: + collect(getattr(cache, "state", ())) + return arrays def async_chained_decode_mlx( self, diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 76f5cc13a..4a6d0df47 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -828,6 +828,11 @@ class SchedulerBatchResultProcessor: else: if self.server_args.enable_hisparse: self.hisparse_coordinator.request_finished(req) + prepare_release = getattr( + self.model_worker, "prepare_for_kv_cache_release", None + ) + if callable(prepare_release): + prepare_release(req) release_kv_cache(req, self.tree_cache) req.time_stats.set_completion_time() diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index 713bf308f..0a8475cb6 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -9,6 +9,9 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.unified_cache_components import ComponentType + from sglang.srt.mem_cache.unified_cache_components.tree_component import ( + TreeComponent, + ) @dataclasses.dataclass @@ -44,3 +47,6 @@ class CacheInitParams: cache_ttl_seconds: Optional[float] = None tree_components: Optional[tuple[ComponentType, ...]] = None + component_registry_override: Optional[dict[ComponentType, type[TreeComponent]]] = ( + None + ) diff --git a/python/sglang/srt/mem_cache/registry.py b/python/sglang/srt/mem_cache/registry.py index c91aae91e..f619f17ad 100644 --- a/python/sglang/srt/mem_cache/registry.py +++ b/python/sglang/srt/mem_cache/registry.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional from sglang.srt.environ import envs from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.utils.tensor_bridge import use_mlx if TYPE_CHECKING: from sglang.srt.configs.model_config import ModelConfig @@ -94,7 +95,7 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: logger.info("Using experimental C++ radix tree implementation.") return RadixCacheCpp(params=params, server_args=server_args) - if envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get(): + if envs.SGLANG_ENABLE_UNIFIED_RADIX_TREE.get() or use_mlx(): from sglang.srt.mem_cache.unified_cache_components import ComponentType from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache @@ -104,6 +105,14 @@ def default_radix_cache_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: ComponentType.SWA if ctx.is_hybrid_swa else ComponentType.MAMBA ) params.tree_components = tuple(tree_components) + if use_mlx() and ctx.is_hybrid_ssm: + from sglang.srt.hardware_backend.mlx.kv_cache.auxiliary_state import ( + MlxAuxiliaryStateComponent, + ) + + params.component_registry_override = { + ComponentType.MAMBA: MlxAuxiliaryStateComponent, + } cache = UnifiedRadixCache(params) if ctx.enable_hierarchical_cache: cache.init_hicache(server_args, params) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 22f418541..93b6924fb 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -258,8 +258,14 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): assert params.tree_components is not None self.tree_components = tuple(params.tree_components) + component_registry = COMPONENT_REGISTRY + if params.component_registry_override: + component_registry = { + **COMPONENT_REGISTRY, + **params.component_registry_override, + } self.components: dict[ComponentType, TreeComponent] = { - ct: COMPONENT_REGISTRY[ct](self, params) for ct in self.tree_components + ct: component_registry[ct](self, params) for ct in self.tree_components } self._components_tuple: tuple[TreeComponent, ...] = tuple( self.components.values() diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py new file mode 100644 index 000000000..e82b1e8dc --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -0,0 +1,1407 @@ +"""Unit tests for MLX attention discovery and generic cache handling.""" + +from __future__ import annotations + +import importlib.util +import unittest +from types import SimpleNamespace + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + +_HAS_MLX = importlib.util.find_spec("mlx") is not None +_SKIP_REASON = "requires mlx" + +if _HAS_MLX: + import mlx.core as mx + import mlx.nn as nn + import torch + from mlx_lm.models.cache import ArraysCache + + import sglang.srt.hardware_backend.mlx.aot as mlx_aot + from sglang.srt.hardware_backend.mlx.aot import ( + MlxAOTKernelSet, + MlxAOTRoPEKernel, + ) + from sglang.srt.hardware_backend.mlx.kv_cache import ( + BatchedDecodeContext, + ContiguousAttentionKVCache, + MlxAttentionKVPool, + MLXAttentionWrapper, + MlxAuxiliaryStateComponent, + MlxAuxiliaryStatePool, + MlxAuxiliaryStateReqToTokenPool, + MlxModelCacheLayout, + find_attention_layers, + is_attention_module, + patch_model_attention, + ) + from sglang.srt.hardware_backend.mlx.model_runner import ( + MlxModelRunner, + MlxPendingDecode, + ) + from sglang.srt.hardware_backend.mlx.scheduler_mixin import ( + MlxPendingJob, + SchedulerMlxOverlapMixin, + ) + from sglang.srt.managers.scheduler_components import ( + batch_result_processor as batch_result_processor_module, + ) + from sglang.srt.managers.scheduler_components.batch_result_processor import ( + SchedulerBatchResultProcessor, + ) + from sglang.srt.managers.utils import GenerationBatchResult + from sglang.srt.mem_cache.base_prefix_cache import InsertParams, InsertResult + from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler + + +def _set_runner_cache_layout( + runner, + *, + num_layers: int, + attention_layer_indices: list[int], + attention_modules: dict[int, object] | None = None, +) -> None: + attention_modules = attention_modules or {} + attention_set = set(attention_layer_indices) + layers = [] + attrs = [] + for layer_idx in range(num_layers): + if layer_idx in attention_set: + attrs.append("self_attn") + layers.append( + SimpleNamespace(self_attn=attention_modules.get(layer_idx, object())) + ) + else: + attrs.append(None) + layers.append(SimpleNamespace()) + runner._cache_layout = MlxModelCacheLayout.from_attention_discovery(layers, attrs) + + +def _set_runner_decode_context_defaults(runner) -> None: + runner._aot_kernels = MlxAOTKernelSet() + runner._attention_kv_pool = None + runner._req_pool_idx = {} + runner._req_to_token_pool = None + + +def _set_dummy_server_args_for_auxiliary_state_tests() -> None: + server_args = ServerArgs(model_path="dummy", page_size=1) + server_args._mamba_cache_chunk_size = 64 + set_global_server_args_for_scheduler(server_args) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestMlxAttentionPatching(unittest.TestCase): + def test_standard_attention_is_patched_once(self): + model = FakeModel( + [ + FakeLayer("self_attn", FakeAttention()), + FakeLayer("self_attn", FakeAttention()), + ] + ) + + layers, attrs = find_attention_layers(model) + + self.assertEqual(len(layers), 2) + self.assertEqual(attrs, ["self_attn", "self_attn"]) + self.assertEqual(patch_model_attention(model), 2) + self.assertIsInstance(model.layers[0].self_attn, MLXAttentionWrapper) + self.assertIsInstance(model.layers[1].self_attn, MLXAttentionWrapper) + self.assertEqual(patch_model_attention(model), 0) + + def test_alias_head_names_are_supported(self): + model = FakeModel([FakeLayer("attention", FakeAttention(use_aliases=True))]) + + _, attrs = find_attention_layers(model) + + self.assertEqual(attrs, ["attention"]) + self.assertEqual(patch_model_attention(model), 1) + self.assertIsInstance(model.layers[0].attention, MLXAttentionWrapper) + + def test_aot_rope_kernel_build_uses_head_aliases(self): + attn = FakeAttention(use_aliases=True) + attn.rope = SimpleNamespace(dims=2, traditional=False, base=10000.0) + original_loader = mlx_aot._load_metal_rope_pool_fused + mlx_aot._load_metal_rope_pool_fused = lambda: object() + try: + kernel = mlx_aot._build_rope_kernel( + mlx_aot.MlxAOTKernelBuildInputs( + sample_attn=attn, + n_kv_heads=1, + head_dim=2, + ) + ) + finally: + mlx_aot._load_metal_rope_pool_fused = original_loader + + self.assertTrue(kernel.enabled) + self.assertEqual(kernel.config["num_qo_heads"], 2) + + def test_auxiliary_state_model_returns_per_layer_attention_attrs(self): + model = FakeModel( + [ + FakeLayer("linear_attn", ProjectionOnlyMixer()), + FakeLayer("self_attn", FakeAttention()), + FakeLayer("linear_attn", ProjectionOnlyMixer()), + ] + ) + + _, attrs = find_attention_layers(model) + + self.assertEqual(attrs, [None, "self_attn", None]) + self.assertEqual(patch_model_attention(model), 1) + self.assertFalse(isinstance(model.layers[0].linear_attn, MLXAttentionWrapper)) + self.assertIsInstance(model.layers[1].self_attn, MLXAttentionWrapper) + + def test_projection_only_mixer_is_not_attention(self): + self.assertFalse(is_attention_module(ProjectionOnlyMixer())) + + def test_cache_layout_separates_attention_and_auxiliary_layers(self): + layout = MlxModelCacheLayout.from_attention_discovery( + [object(), object(), object(), object()], + [None, "self_attn", None, "self_attn"], + ) + + self.assertEqual(layout.num_layers, 4) + self.assertEqual(layout.attention_layer_indices, (1, 3)) + self.assertEqual(layout.auxiliary_layer_indices, (0, 2)) + self.assertEqual(layout.attention_pool_index(1), 0) + self.assertEqual(layout.attention_pool_index(3), 1) + self.assertTrue(layout.has_auxiliary_state) + + def test_gated_query_projection_keeps_attention_width(self): + inner = FakeGatedAttention() + wrapper = MLXAttentionWrapper(inner, layer_idx=0) + cache = ContiguousAttentionKVCache( + n_kv_heads=1, head_dim=2, max_seq_len=4, dtype=mx.float32 + ) + ctx = BatchedDecodeContext( + batch_size=1, + seq_lens=[0], + attention_layer_caches=[[cache]], + ) + + out = wrapper._batched_decode(mx.zeros((1, 1, 4), dtype=mx.float32), ctx) + mx.eval(out) + + self.assertEqual(out.shape, (1, 1, 4)) + self.assertEqual(inner.o_proj.last_input_shape, (1, 1, 4)) + + def test_attn_config_uses_float_dtype_for_quantized_projection(self): + runner = object.__new__(MlxModelRunner) + attn = FakeAttention() + attn.k_proj.weight = mx.zeros((2, 4), dtype=mx.uint32) + _set_runner_cache_layout( + runner, + num_layers=1, + attention_layer_indices=[0], + attention_modules={0: attn}, + ) + + n_kv_heads, head_dim, dtype = MlxModelRunner._get_attn_config(runner) + + self.assertEqual(n_kv_heads, 1) + self.assertEqual(head_dim, 2) + self.assertEqual(dtype, mx.float32) + + def test_attn_config_rejects_heterogeneous_kv_shapes(self): + runner = object.__new__(MlxModelRunner) + first = FakeAttention() + second = FakeAttention() + second.n_kv_heads = 2 + _set_runner_cache_layout( + runner, + num_layers=2, + attention_layer_indices=[0, 1], + attention_modules={0: first, 1: second}, + ) + + with self.assertRaisesRegex( + NotImplementedError, + "uniform softmax-attention KV shape", + ): + MlxModelRunner._get_attn_config(runner) + + def test_attn_config_rejects_sliding_window_attention(self): + runner = object.__new__(MlxModelRunner) + _set_runner_cache_layout( + runner, + num_layers=1, + attention_layer_indices=[0], + attention_modules={0: FakeAttention()}, + ) + runner._cache_layout.layers[0].use_sliding = True + + with self.assertRaisesRegex( + NotImplementedError, + "sliding-window attention", + ): + MlxModelRunner._get_attn_config(runner) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase): + def test_dense_prefill_keeps_pool_backed_radix_path(self): + runner = object.__new__(MlxModelRunner) + runner.model = FakeDenseModel(num_layers=2) + _set_runner_cache_layout( + runner, + num_layers=2, + attention_layer_indices=[0, 1], + ) + runner._max_seq_len = 8 + runner._cache_pool = [] + runner.disable_radix_cache = False + runner._attention_kv_pool = MlxAttentionKVPool( + pool_size=8, + num_layers=2, + n_kv_heads=1, + head_dim=2, + dtype=mx.float32, + ) + runner._req_to_token_pool = None + runner._req_caches = {} + runner._req_token_ids = {} + runner._req_pool_idx = {} + runner._req_synced_offset = {} + prefix_slots = mx.array([2, 3], dtype=mx.int32) + k_prefix = mx.stack( + [ + mx.ones((2, 1, 2), dtype=mx.float32) * 10, + mx.ones((2, 1, 2), dtype=mx.float32) * 20, + ] + ) + runner._attention_kv_pool.set_kv_all_layers( + prefix_slots, k_prefix, k_prefix * 2 + ) + mx.eval(*runner._attention_kv_pool.all_buffers()) + + pending = runner.prefill_start( + req_id="r0", + new_token_ids=[13], + full_token_ids=[11, 12, 13], + prefix_slot_ids=[2, 3], + new_slot_ids=[4], + req_pool_idx=0, + ) + MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache) + mx.eval(*runner._attention_kv_pool.all_buffers()) + runner.prefill_finalize(pending) + + self.assertEqual(runner.model.seen_inputs, [[[13]]]) + self.assertEqual(runner.model.seen_offsets, [[2, 2]]) + self.assertEqual(pending.synced_offset, 3) + self.assertTrue( + all(isinstance(c, ContiguousAttentionKVCache) for c in pending.cache) + ) + layer0_k, layer0_v = runner._attention_kv_pool.get_kv( + 0, mx.array([4], dtype=mx.int32) + ) + layer1_k, layer1_v = runner._attention_kv_pool.get_kv( + 1, mx.array([4], dtype=mx.int32) + ) + mx.eval(layer0_k, layer0_v, layer1_k, layer1_v) + self.assertEqual(layer0_k.tolist(), [[[1.0, 1.0]]]) + self.assertEqual(layer0_v.tolist(), [[[2.0, 2.0]]]) + self.assertEqual(layer1_k.tolist(), [[[2.0, 2.0]]]) + self.assertEqual(layer1_v.tolist(), [[[4.0, 4.0]]]) + + def test_dense_decode_uses_batched_attention_for_single_and_multi_request(self): + for req_ids in (["r0"], ["r0", "r1"]): + with self.subTest(req_ids=req_ids): + runner = object.__new__(MlxModelRunner) + _set_runner_cache_layout( + runner, + num_layers=1, + attention_layer_indices=[0], + ) + runner._req_caches = {rid: [object()] for rid in req_ids} + runner._req_token_ids = { + rid: [idx + 10] for idx, rid in enumerate(req_ids) + } + calls = [] + + def fake_batched(caches, batched_input, helper_req_ids): + calls.append( + (len(caches), batched_input.tolist(), list(helper_req_ids)) + ) + return mx.array(list(range(len(caches))), dtype=mx.int32) + + def fail_native(*args, **kwargs): + raise AssertionError("dense decode should use batched attention") + + runner._decode_with_batched_attention = fake_batched + runner._decode_with_native_cache = fail_native + + pending = runner.decode_batch_start(req_ids) + + self.assertEqual( + calls, + [ + ( + len(req_ids), + [[idx + 10] for idx in range(len(req_ids))], + req_ids, + ) + ], + ) + self.assertEqual( + pending.lazy_tokens.tolist(), + list(range(len(req_ids))), + ) + + def test_dense_chained_decode_uses_batched_attention_for_single_request(self): + runner = object.__new__(MlxModelRunner) + _set_runner_cache_layout( + runner, + num_layers=1, + attention_layer_indices=[0], + ) + calls = [] + + def fake_batched(caches, batched_input, helper_req_ids): + calls.append((len(caches), batched_input.tolist(), list(helper_req_ids))) + return mx.array([8], dtype=mx.int32) + + def fail_native(*args, **kwargs): + raise AssertionError("dense chained decode should use batched attention") + + runner._decode_with_batched_attention = fake_batched + runner._decode_with_native_cache = fail_native + prev = MlxPendingDecode( + lazy_tokens=mx.array([7], dtype=mx.int32), + req_ids=["r0"], + caches=[[object()]], + ) + + pending = runner.decode_batch_start_chained(prev) + + self.assertEqual(calls, [(1, [[7]], ["r0"])]) + self.assertEqual(pending.lazy_tokens.tolist(), [8]) + + def test_decode_finalize_does_not_snapshot_auxiliary_state(self): + runner = object.__new__(MlxModelRunner) + runner._req_token_ids = {"r0": [8]} + runner._decode_step_ct = 0 + calls = [] + runner._store_auxiliary_state = lambda req_pool_idx, cache: calls.append( + (req_pool_idx, cache) + ) + pending = MlxPendingDecode( + lazy_tokens=mx.array([9], dtype=mx.int32), + req_ids=["r0"], + caches=[[object()]], + ) + + next_tokens = runner.decode_batch_finalize(pending) + + self.assertEqual(next_tokens, [9]) + self.assertEqual(runner._req_token_ids["r0"], [8, 9]) + self.assertEqual(calls, []) + + def test_store_auxiliary_state_for_request_snapshots_on_demand(self): + runner = object.__new__(MlxModelRunner) + cache = [object()] + runner._req_pool_idx = {"r0": 3} + runner._req_caches = {"r0": cache} + calls = [] + runner._store_auxiliary_state = lambda req_pool_idx, cache_arg: calls.append( + (req_pool_idx, cache_arg) + ) + + runner.store_auxiliary_state_for_request("r0") + runner.store_auxiliary_state_for_request("missing") + + self.assertEqual(calls, [(3, cache)]) + + def test_dense_batched_attention_helper_supports_single_request(self): + runner = object.__new__(MlxModelRunner) + model = FakeWrappedAttentionModel() + runner.model = model + _set_runner_cache_layout( + runner, + num_layers=1, + attention_layer_indices=[0], + ) + _set_runner_decode_context_defaults(runner) + cache = [ + [ + ContiguousAttentionKVCache( + n_kv_heads=1, + head_dim=2, + max_seq_len=4, + dtype=mx.float32, + ) + ] + ] + + lazy_tokens = runner._decode_with_batched_attention( + cache, + mx.array([[7]], dtype=mx.int32), + ["r0"], + ) + mx.eval(lazy_tokens, *MlxModelRunner._cache_state_arrays(cache)) + + self.assertEqual(lazy_tokens.tolist(), [0]) + self.assertEqual(cache[0][0].offset, 1) + self.assertEqual(model.seen_inputs, [[[7]]]) + self.assertEqual(model.seen_cache_types, [["AttentionOffsetCache"]]) + + def test_batched_decode_context_resolves_aot_rope_slots_from_request_ids(self): + cache0 = ContiguousAttentionKVCache( + n_kv_heads=1, + head_dim=2, + max_seq_len=4, + dtype=mx.float32, + ) + cache1 = ContiguousAttentionKVCache( + n_kv_heads=1, + head_dim=2, + max_seq_len=4, + dtype=mx.float32, + ) + cache0.offset = 1 + cache1.offset = 2 + kernel_set = MlxAOTKernelSet( + rope=MlxAOTRoPEKernel( + base=10000.0, + config={ + "head_dim": 2, + "num_qo_heads": 1, + "num_kv_heads": 1, + }, + rope_pool_fused=object(), + ) + ) + req_to_token_pool = SimpleNamespace( + req_to_token=torch.tensor( + [ + [0, 41, 42], + [0, 51, 52], + ], + dtype=torch.int64, + ) + ) + + ctx = BatchedDecodeContext.from_decode( + caches=[[cache0], [cache1]], + req_ids=["r0", "r1"], + aot_kernels=kernel_set, + kv_pool=object(), + req_pool_idx={"r0": 0, "r1": 1}, + req_to_token_pool=req_to_token_pool, + attention_layer_indices=[0], + ) + + self.assertEqual(ctx.seq_lens, [1, 2]) + self.assertIsNotNone(ctx.aot.rope) + self.assertEqual(ctx.aot.rope.new_token_slots.tolist(), [41, 52]) + + def test_auxiliary_decode_uses_hybrid_batching_for_multi_request(self): + runner = object.__new__(MlxModelRunner) + _set_runner_cache_layout( + runner, + num_layers=2, + attention_layer_indices=[1], + ) + req_ids = ["r0", "r1"] + runner._req_caches = {rid: [object(), object()] for rid in req_ids} + runner._req_token_ids = {rid: [idx + 20] for idx, rid in enumerate(req_ids)} + calls = [] + + def fake_hybrid(caches, batched_input, helper_req_ids): + calls.append((len(caches), batched_input.tolist(), list(helper_req_ids))) + return mx.array([4, 5], dtype=mx.int32) + + def fail_batched(*args, **kwargs): + raise AssertionError( + "auxiliary decode should use hybrid batching, not full batched" + ) + + runner._decode_with_hybrid_batching = fake_hybrid + runner._decode_with_batched_attention = fail_batched + + pending = runner.decode_batch_start(req_ids) + + self.assertEqual(calls, [(2, [[20], [21]], req_ids)]) + self.assertEqual(pending.lazy_tokens.tolist(), [4, 5]) + + def test_auxiliary_layer_batches_mergeable_native_cache(self): + runner = object.__new__(MlxModelRunner) + layer = FakeBatchableAuxiliaryLayer() + cache0 = ArraysCache(size=1) + cache1 = ArraysCache(size=1) + + out = runner._decode_auxiliary_layer( + layer, + mx.zeros((2, 1, 4), dtype=mx.float32), + [cache0, cache1], + ) + mx.eval(out, cache0[0], cache1[0]) + + self.assertEqual(layer.input_layernorm.seen_shapes, [(2, 1, 4)]) + self.assertEqual(layer.linear_attn.seen_shapes, [(2, 1, 4)]) + self.assertEqual(layer.post_attention_layernorm.seen_shapes, [(2, 1, 4)]) + self.assertEqual(layer.mlp.seen_shapes, [(2, 1, 4)]) + self.assertEqual(layer.linear_attn.cache_type, "ArraysCache") + self.assertEqual( + out.tolist(), + [[[2.0, 2.0, 2.0, 2.0]], [[2.0, 2.0, 2.0, 2.0]]], + ) + self.assertEqual(cache0[0].tolist(), [[0.0]]) + self.assertEqual(cache1[0].tolist(), [[1.0]]) + + def test_arrays_cache_auxiliary_batching_uses_fast_merge(self): + runner = object.__new__(MlxModelRunner) + layer = FakeBatchableAuxiliaryLayer() + cache0 = ArraysCache(size=1) + cache1 = ArraysCache(size=1) + original_merge = ArraysCache.merge + + def fail_merge(cls, caches): + raise AssertionError("ArraysCache fast path should not call merge()") + + ArraysCache.merge = classmethod(fail_merge) + try: + out = runner._decode_auxiliary_layer( + layer, + mx.zeros((2, 1, 4), dtype=mx.float32), + [cache0, cache1], + ) + mx.eval(out, cache0[0], cache1[0]) + finally: + ArraysCache.merge = original_merge + + self.assertEqual(layer.linear_attn.cache_type, "ArraysCache") + self.assertEqual( + out.tolist(), + [[[2.0, 2.0, 2.0, 2.0]], [[2.0, 2.0, 2.0, 2.0]]], + ) + self.assertEqual(cache0[0].tolist(), [[0.0]]) + self.assertEqual(cache1[0].tolist(), [[1.0]]) + + def test_auxiliary_layer_split_back_copies_cache_metadata(self): + runner = object.__new__(MlxModelRunner) + layer = FakeBatchableAuxiliaryLayer() + cache0 = FakeMergeableAuxiliaryCache(tag="old0") + cache1 = FakeMergeableAuxiliaryCache(tag="old1") + + out = runner._decode_auxiliary_layer( + layer, + mx.zeros((2, 1, 4), dtype=mx.float32), + [cache0, cache1], + ) + mx.eval(out, cache0[0], cache1[0]) + + self.assertEqual(layer.linear_attn.cache_type, "FakeMergeableAuxiliaryCache") + self.assertEqual(cache0.tag, "split-0") + self.assertEqual(cache1.tag, "split-1") + self.assertEqual(cache0.extra_metadata, {"idx": 0}) + self.assertEqual(cache1.extra_metadata, {"idx": 1}) + self.assertEqual(cache0[0].tolist(), [[0.0]]) + self.assertEqual(cache1[0].tolist(), [[1.0]]) + + def test_auxiliary_state_prefill_restores_prefix_state(self): + runner = object.__new__(MlxModelRunner) + runner.model = FakeAuxiliaryStateModel() + _set_runner_cache_layout( + runner, + num_layers=2, + attention_layer_indices=[1], + ) + runner._max_seq_len = 8 + runner._cache_pool = [] + runner.disable_radix_cache = False + runner._attention_kv_pool = MlxAttentionKVPool( + pool_size=8, + num_layers=1, + n_kv_heads=1, + head_dim=2, + dtype=mx.float32, + ) + runner._req_to_token_pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + runner._req_caches = {} + runner._req_token_ids = {} + runner._req_pool_idx = {} + runner._req_synced_offset = {} + req = FakeRequest() + runner._req_to_token_pool.alloc([req]) + runner._req_to_token_pool.auxiliary_state_pool.store_cache( + req.mamba_pool_idx, + [FakeNativeCache(mx.array([42.0], dtype=mx.float32)), None], + [0], + ) + + pending = runner.prefill_start( + req_id="r0", + new_token_ids=[13], + full_token_ids=[11, 12, 13], + prefix_slot_ids=[2, 3], + new_slot_ids=[4], + req_pool_idx=req.req_pool_idx, + ) + MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache) + runner.prefill_finalize(pending) + + self.assertEqual(runner.model.seen_inputs, [[[13]]]) + self.assertEqual(runner.model.seen_auxiliary_states, [[42.0]]) + self.assertEqual(pending.synced_offset, 3) + self.assertIsInstance(pending.cache[0], FakeNativeCache) + self.assertIsInstance(pending.cache[1], ContiguousAttentionKVCache) + restored = [FakeNativeCache(), None] + runner._req_to_token_pool.auxiliary_state_pool.restore_cache( + req.mamba_pool_idx, restored, [0] + ) + self.assertEqual(restored[0].state[0].tolist(), [1.0]) + + def test_auxiliary_state_prefill_tracks_chunk_aligned_auxiliary_state(self): + _set_dummy_server_args_for_auxiliary_state_tests() + runner = object.__new__(MlxModelRunner) + runner.model = FakeAuxiliaryStateModel() + _set_runner_cache_layout( + runner, + num_layers=2, + attention_layer_indices=[1], + ) + runner._max_seq_len = 128 + runner._cache_pool = [] + runner.disable_radix_cache = False + runner._attention_kv_pool = MlxAttentionKVPool( + pool_size=96, + num_layers=1, + n_kv_heads=1, + head_dim=2, + dtype=mx.float32, + ) + runner._req_to_token_pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=128, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + runner._req_caches = {} + runner._req_token_ids = {} + runner._req_pool_idx = {} + runner._req_synced_offset = {} + req = FakeRequest() + runner._req_to_token_pool.alloc([req]) + token_ids = list(range(70)) + + pending = runner.prefill_start( + req_id="r0", + new_token_ids=token_ids, + full_token_ids=token_ids, + prefix_slot_ids=[], + new_slot_ids=list(range(1, 71)), + req_pool_idx=req.req_pool_idx, + req=req, + ) + MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache) + runner.prefill_finalize(pending) + tracked = [FakeNativeCache(), None] + runner._req_to_token_pool.auxiliary_state_pool.restore_cache( + req.mamba_ping_pong_track_buffer[0], tracked, [0] + ) + + self.assertEqual([len(x[0]) for x in runner.model.seen_inputs], [64, 6]) + self.assertEqual(req.mamba_last_track_seqlen, 64) + self.assertEqual(tracked[0].state[0].tolist(), [64.0]) + self.assertEqual(pending.synced_offset, 70) + + def test_auxiliary_state_prefill_advances_tracked_boundary_after_cached_prefix( + self, + ): + _set_dummy_server_args_for_auxiliary_state_tests() + runner = object.__new__(MlxModelRunner) + runner.model = FakeAuxiliaryStateModel() + _set_runner_cache_layout( + runner, + num_layers=2, + attention_layer_indices=[1], + ) + runner._max_seq_len = 512 + runner._cache_pool = [] + runner.disable_radix_cache = False + runner._attention_kv_pool = MlxAttentionKVPool( + pool_size=320, + num_layers=1, + n_kv_heads=1, + head_dim=2, + dtype=mx.float32, + ) + runner._req_to_token_pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=512, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + runner._req_caches = {} + runner._req_token_ids = {} + runner._req_pool_idx = {} + runner._req_synced_offset = {} + req = FakeRequest() + runner._req_to_token_pool.alloc([req]) + runner._req_to_token_pool.auxiliary_state_pool.store_cache( + req.mamba_pool_idx, + [FakeNativeCache(mx.array([64.0], dtype=mx.float32)), None], + [0], + ) + token_ids = list(range(257)) + + pending = runner.prefill_start( + req_id="r0", + new_token_ids=token_ids[64:], + full_token_ids=token_ids, + prefix_slot_ids=list(range(1, 65)), + new_slot_ids=list(range(65, 258)), + req_pool_idx=req.req_pool_idx, + req=req, + ) + MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache) + runner.prefill_finalize(pending) + tracked = [FakeNativeCache(), None] + runner._req_to_token_pool.auxiliary_state_pool.restore_cache( + req.mamba_ping_pong_track_buffer[0], tracked, [0] + ) + + self.assertEqual([len(x[0]) for x in runner.model.seen_inputs], [192, 1]) + self.assertEqual(runner.model.seen_auxiliary_states, [[64.0], [192.0]]) + self.assertEqual(req.mamba_last_track_seqlen, 256) + self.assertEqual(tracked[0].state[0].tolist(), [192.0]) + self.assertEqual(pending.synced_offset, 257) + + def test_cache_arrays_flattens_native_array_cache_state(self): + cache = FakeNestedStateCache() + + arrays = MlxModelRunner._cache_arrays(cache) + + self.assertEqual(len(arrays), 2) + self.assertTrue(all(isinstance(arr, mx.array) for arr in arrays)) + + def test_auxiliary_state_pool_tracks_scheduler_slots_and_snapshots(self): + pool = MlxAuxiliaryStatePool(size=4, device="cpu") + + first = pool.alloc(2) + cache = [FakeNativeCache(mx.array([1.0], dtype=mx.float32))] + pool.store_cache(first[0], cache, [0]) + cache[0].state[0][0] = 9.0 + forked = pool.fork_from(first[0].unsqueeze(0)) + restored = [FakeNativeCache()] + pool.restore_cache(forked[0], restored, [0]) + pool.free(first) + + self.assertEqual(first.tolist(), [1, 2]) + self.assertEqual(forked.tolist(), [3]) + self.assertEqual(restored[0].state[0].tolist(), [1.0]) + self.assertEqual(pool.available_size(), 3) + + def test_auxiliary_state_pool_restores_instance_meta_state(self): + pool = MlxAuxiliaryStatePool(size=2, device="cpu") + slot = pool.alloc(1) + cache = [ + FakeNativeCache( + mx.array([1.0], dtype=mx.float32), + meta_state={"seen": mx.array([3.0], dtype=mx.float32)}, + ) + ] + pool.store_cache(slot[0], cache, [0]) + cache[0].meta_state["seen"][0] = 9.0 + + restored = [FakeNativeCache(meta_state={})] + pool.restore_cache(slot[0], restored, [0]) + + self.assertEqual(restored[0].meta_state["seen"].tolist(), [3.0]) + + def test_auxiliary_state_req_pool_maps_request_indices(self): + pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + req = FakeRequest() + + req_indices = pool.alloc([req]) + auxiliary_state_idx = pool.get_auxiliary_state_indices(req.req_pool_idx) + pool.free(req) + + self.assertEqual(req_indices, [1]) + self.assertIsNotNone(auxiliary_state_idx) + self.assertIsNone(req.req_pool_idx) + self.assertIsNotNone(req.mamba_pool_idx) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) + pool.free_auxiliary_state_cache(req) + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual(pool.available_size(), 2) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 4) + + def test_auxiliary_state_req_pool_can_keep_tracked_auxiliary_slot(self): + pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + req = FakeRequest() + pool.alloc([req]) + req.mamba_ping_pong_track_buffer = pool.auxiliary_state_pool.alloc(1) + req.mamba_next_track_idx = 0 + + pool.free_auxiliary_state_cache(req, track_buffer_to_keep=0) + + self.assertIsNone(req.mamba_pool_idx) + self.assertIsNone(req.mamba_ping_pong_track_buffer) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) + + def test_auxiliary_state_component_inserts_tracked_slot_and_frees_live_slot(self): + pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + req = FakeRequest() + pool.alloc([req]) + req.mamba_ping_pong_track_buffer = pool.auxiliary_state_pool.alloc(1) + req.mamba_next_track_idx = 0 + req.mamba_last_track_seqlen = 64 + component = MlxAuxiliaryStateComponent( + SimpleNamespace(req_to_token_pool=pool), + SimpleNamespace(enable_mamba_extra_buffer=False), + ) + insert_params = InsertParams() + + cache_len = component.prepare_for_caching_req( + req=req, + insert_params=insert_params, + token_ids_len=70, + is_finished=True, + ) + component.cleanup_after_caching_req( + req=req, + is_finished=True, + insert_result=InsertResult(prefix_len=0, mamba_exist=False), + insert_params=insert_params, + ) + + self.assertEqual(cache_len, 64) + self.assertTrue(getattr(insert_params, "mlx_auxiliary_state_uses_track_slot")) + self.assertEqual(insert_params.mamba_value.tolist(), [2]) + self.assertIsNone(req.mamba_pool_idx) + self.assertIsNone(req.mamba_ping_pong_track_buffer) + self.assertIsNone(req.mamba_last_track_seqlen) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) + + def test_auxiliary_state_component_unfinished_frees_tracked_source_slot(self): + pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + req = FakeRequest() + pool.alloc([req]) + req.mamba_ping_pong_track_buffer = pool.auxiliary_state_pool.alloc(1) + req.mamba_next_track_idx = 0 + req.mamba_last_track_seqlen = 64 + component = MlxAuxiliaryStateComponent( + SimpleNamespace(req_to_token_pool=pool), + SimpleNamespace(enable_mamba_extra_buffer=False), + ) + insert_params = InsertParams() + + cache_len = component.prepare_for_caching_req( + req=req, + insert_params=insert_params, + token_ids_len=70, + is_finished=False, + ) + component.cleanup_after_caching_req( + req=req, + is_finished=False, + insert_result=InsertResult(prefix_len=0, mamba_exist=False), + insert_params=insert_params, + ) + + self.assertEqual(cache_len, 64) + self.assertEqual(insert_params.mamba_value.tolist(), [3]) + self.assertIsNotNone(req.mamba_pool_idx) + self.assertIsNone(req.mamba_ping_pong_track_buffer) + self.assertIsNone(req.mamba_last_track_seqlen) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 2) + + def test_auxiliary_state_component_keeps_new_live_slot_owned_by_radix(self): + pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + req = FakeRequest() + pool.alloc([req]) + component = MlxAuxiliaryStateComponent( + SimpleNamespace(req_to_token_pool=pool), + SimpleNamespace(enable_mamba_extra_buffer=False), + ) + insert_params = InsertParams() + + cache_len = component.prepare_for_caching_req( + req=req, + insert_params=insert_params, + token_ids_len=7, + is_finished=True, + ) + component.cleanup_after_caching_req( + req=req, + is_finished=True, + insert_result=InsertResult(prefix_len=0, mamba_exist=False), + insert_params=insert_params, + ) + + self.assertEqual(cache_len, 7) + self.assertFalse(getattr(insert_params, "mlx_auxiliary_state_uses_track_slot")) + self.assertEqual(insert_params.mamba_value.tolist(), [1]) + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) + + def test_auxiliary_state_component_frees_stale_track_slot_when_live_slot_inserted( + self, + ): + pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + req = FakeRequest() + pool.alloc([req]) + req.mamba_ping_pong_track_buffer = pool.auxiliary_state_pool.alloc(1) + req.mamba_next_track_idx = 0 + component = MlxAuxiliaryStateComponent( + SimpleNamespace(req_to_token_pool=pool), + SimpleNamespace(enable_mamba_extra_buffer=False), + ) + insert_params = InsertParams() + + cache_len = component.prepare_for_caching_req( + req=req, + insert_params=insert_params, + token_ids_len=7, + is_finished=True, + ) + component.cleanup_after_caching_req( + req=req, + is_finished=True, + insert_result=InsertResult(prefix_len=0, mamba_exist=False), + insert_params=insert_params, + ) + + self.assertEqual(cache_len, 7) + self.assertFalse(getattr(insert_params, "mlx_auxiliary_state_uses_track_slot")) + self.assertEqual(insert_params.mamba_value.tolist(), [1]) + self.assertIsNone(req.mamba_pool_idx) + self.assertIsNone(req.mamba_ping_pong_track_buffer) + self.assertIsNone(req.mamba_next_track_idx) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 3) + + def test_auxiliary_state_component_frees_duplicate_live_slot(self): + pool = MlxAuxiliaryStateReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + auxiliary_state_size=4, + ) + req = FakeRequest() + pool.alloc([req]) + component = MlxAuxiliaryStateComponent( + SimpleNamespace(req_to_token_pool=pool), + SimpleNamespace(enable_mamba_extra_buffer=False), + ) + insert_params = InsertParams() + + component.prepare_for_caching_req( + req=req, + insert_params=insert_params, + token_ids_len=7, + is_finished=True, + ) + component.cleanup_after_caching_req( + req=req, + is_finished=True, + insert_result=InsertResult(prefix_len=7, mamba_exist=True), + insert_params=insert_params, + ) + + self.assertIsNone(req.mamba_pool_idx) + self.assertEqual(pool.auxiliary_state_pool.available_size(), 4) + + +@unittest.skipUnless(_HAS_MLX, _SKIP_REASON) +class TestMlxOverlapScheduler(unittest.TestCase): + def test_finalize_pending_job_updates_scheduler_last_batch(self): + token_ids = torch.tensor([7], dtype=torch.long) + scheduler = FakeOverlapScheduler(token_ids) + stale_batch = SimpleNamespace(input_ids=None) + batch_copy = SimpleNamespace(input_ids=None) + schedule_batch = SimpleNamespace(input_ids=None) + scheduler.last_batch = stale_batch + + pending = MlxPendingJob( + lazy_tokens=None, + prefills=["prefill"], + extends=[], + decode=None, + mode="extend", + batch_copy=batch_copy, + schedule_batch=schedule_batch, + reqs=[SimpleNamespace(rid="r0")], + ) + + scheduler._finalize_mlx_pending_job(pending) + + self.assertIs(scheduler.last_batch, schedule_batch) + self.assertTrue(torch.equal(batch_copy.input_ids, token_ids)) + self.assertTrue(torch.equal(schedule_batch.input_ids, token_ids)) + self.assertIs(scheduler.processed_batch, batch_copy) + self.assertIs(scheduler.processed_result, scheduler.tp_worker.result) + + def test_finished_request_snapshots_before_release(self): + events = [] + tree_cache = object() + processor = SchedulerBatchResultProcessor( + is_generation=True, + disaggregation_mode=None, + enable_overlap=False, + enable_overlap_mlx=False, + server_args=SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False, + enable_hisparse=False, + ), + model_config=None, + token_to_kv_pool_allocator=None, + tree_cache=tree_cache, + hisparse_coordinator=None, + req_to_token_pool=None, + decode_offload_manager=None, + metrics_collector=None, + metrics_reporter=None, + draft_worker=None, + model_worker=SimpleNamespace( + prepare_for_kv_cache_release=lambda req: events.append( + ("prepare", req.rid) + ) + ), + logprob_result_processor=None, + output_streamer=None, + abort_request=lambda req: None, + ) + req = SimpleNamespace( + rid="r0", + finished=lambda: True, + multimodal_inputs=None, + session=None, + return_routed_experts=False, + time_stats=SimpleNamespace( + set_completion_time=lambda: events.append(("completion", "r0")) + ), + ) + original_release = batch_result_processor_module.release_kv_cache + original_get_indexer = batch_result_processor_module.get_global_indexer_capturer + + def fake_release_kv_cache(release_req, tree_cache): + events.append(("release", release_req.rid)) + self.assertIs(tree_cache, processor.tree_cache) + + batch_result_processor_module.release_kv_cache = fake_release_kv_cache + batch_result_processor_module.get_global_indexer_capturer = lambda: None + try: + SchedulerBatchResultProcessor._handle_finished_req( + processor, req, 0, SimpleNamespace(customized_info=None) + ) + finally: + batch_result_processor_module.release_kv_cache = original_release + batch_result_processor_module.get_global_indexer_capturer = ( + original_get_indexer + ) + + self.assertEqual( + events, + [ + ("prepare", "r0"), + ("release", "r0"), + ("completion", "r0"), + ], + ) + + +if _HAS_MLX: + + class FakeProjection(nn.Module): + def __init__(self, out_dim: int = 4): + super().__init__() + self.weight = mx.zeros((out_dim, 4), dtype=mx.float32) + + def __call__(self, x): + shape = (*x.shape[:-1], self.weight.shape[0]) + return mx.zeros(shape, dtype=x.dtype) + + class FakeAttention(nn.Module): + def __init__(self, use_aliases: bool = False): + super().__init__() + if use_aliases: + self.num_attention_heads = 2 + self.num_key_value_heads = 1 + else: + self.n_heads = 2 + self.n_kv_heads = 1 + self.head_dim = 2 + self.scale = self.head_dim**-0.5 + self.q_proj = FakeProjection(4) + self.k_proj = FakeProjection(2) + self.v_proj = FakeProjection(2) + self.o_proj = FakeProjection(4) + self.rope = lambda x, offset=None: x + + class ProjectionOnlyMixer(nn.Module): + def __init__(self): + super().__init__() + self.n_heads = 2 + self.n_kv_heads = 1 + self.q_proj = FakeProjection(4) + self.k_proj = FakeProjection(2) + self.v_proj = FakeProjection(2) + self.o_proj = FakeProjection(4) + + class FakeLayer(nn.Module): + def __init__(self, attr_name: str, module: nn.Module): + super().__init__() + setattr(self, attr_name, module) + + class FakeModel(nn.Module): + def __init__(self, layers): + super().__init__() + self.layers = layers + + class IdentityNorm(nn.Module): + def __call__(self, x): + return x + + class IdentityRope: + def __call__(self, x, offset=None): + return x + + class CapturingOutput(nn.Module): + def __init__(self): + super().__init__() + self.last_input_shape = None + + def __call__(self, x): + self.last_input_shape = x.shape + return x + + class FakeGatedAttention(nn.Module): + def __init__(self): + super().__init__() + self.num_attention_heads = 2 + self.num_key_value_heads = 1 + self.head_dim = 2 + self.scale = self.head_dim**-0.5 + self.q_proj = FakeProjection(8) + self.k_proj = FakeProjection(2) + self.v_proj = FakeProjection(2) + self.o_proj = CapturingOutput() + self.q_norm = IdentityNorm() + self.k_norm = IdentityNorm() + self.rope = IdentityRope() + + class RecordingIdentity(nn.Module): + def __init__(self): + super().__init__() + self.seen_shapes = [] + + def __call__(self, x): + self.seen_shapes.append(x.shape) + return x + + class FakeMergeableLinearAttention(nn.Module): + def __init__(self): + super().__init__() + self.seen_shapes = [] + self.cache_type = None + + def __call__(self, x, mask=None, cache=None): + self.seen_shapes.append(x.shape) + self.cache_type = type(cache).__name__ + cache[0] = mx.arange(x.shape[0], dtype=mx.float32).reshape(x.shape[0], 1) + return x + 1 + + class FakeBatchableAuxiliaryLayer(nn.Module): + def __init__(self): + super().__init__() + self.is_linear = True + self.input_layernorm = RecordingIdentity() + self.linear_attn = FakeMergeableLinearAttention() + self.post_attention_layernorm = RecordingIdentity() + self.mlp = RecordingIdentity() + + class FakeMergeableAuxiliaryCache: + def __init__(self, state=None, tag="init", extra_metadata=None): + self.cache = [state] + self.tag = tag + self.extra_metadata = extra_metadata or {} + + def __getitem__(self, idx): + return self.cache[idx] + + def __setitem__(self, idx, value): + self.cache[idx] = value + + @classmethod + def merge(cls, caches): + merged = cls(tag="merged") + values = [cache[0] for cache in caches] + if all(value is None for value in values): + return merged + merged[0] = mx.concatenate( + [ + ( + value + if value is not None + else mx.zeros_like(next(v for v in values if v is not None)) + ) + for value in values + ], + axis=0, + ) + return merged + + def extract(self, idx): + return type(self)( + self.cache[0][idx : idx + 1], + tag=f"split-{idx}", + extra_metadata={"idx": idx}, + ) + + class FakeNativeCache: + def __init__(self, value=None, meta_state=None): + self._state = [ + value if value is not None else mx.array([0.0], dtype=mx.float32) + ] + if meta_state is not None: + self.meta_state = meta_state + self.lengths = None + self.left_padding = None + + @property + def state(self): + return self._state + + @state.setter + def state(self, value): + self._state = value + + class FakeAuxiliaryStateModel: + def __init__(self): + self.seen_inputs = [] + self.seen_auxiliary_states = [] + + def make_cache(self): + return [FakeNativeCache(), FakeNativeCache()] + + def __call__(self, inputs, cache=None): + self.seen_inputs.append(inputs.tolist()) + if cache is not None: + self.seen_auxiliary_states.append(cache[0].state[0].tolist()) + cache[0].state = [mx.array([float(inputs.shape[1])], dtype=mx.float32)] + keys = mx.ones((1, 1, inputs.shape[1], 2), dtype=mx.float32) + values = keys * 2 + cache[1].update_and_fetch(keys, values) + return mx.zeros((1, inputs.shape[1], 4), dtype=mx.float32) + + class FakeDenseModel: + def __init__(self, num_layers): + self.num_layers = num_layers + self.seen_inputs = [] + self.seen_offsets = [] + + def __call__(self, inputs, cache=None): + self.seen_inputs.append(inputs.tolist()) + if cache is not None: + offsets = [] + for layer_idx in range(self.num_layers): + offsets.append(cache[layer_idx].offset) + scale = float(layer_idx + 1) + keys = mx.ones((1, 1, inputs.shape[1], 2), dtype=mx.float32) * scale + cache[layer_idx].update_and_fetch(keys, keys * 2) + self.seen_offsets.append(offsets) + return mx.zeros((1, inputs.shape[1], 4), dtype=mx.float32) + + class FakeWrappedAttentionModel: + def __init__(self): + self.attn = MLXAttentionWrapper(FakeAttention(), layer_idx=0) + self.seen_inputs = [] + self.seen_cache_types = [] + + def __call__(self, inputs, cache=None): + self.seen_inputs.append(inputs.tolist()) + self.seen_cache_types.append([type(c).__name__ for c in cache]) + hidden = mx.zeros((*inputs.shape, 4), dtype=mx.float32) + self.attn(hidden, cache=cache[0]) + return mx.zeros((*inputs.shape, 8), dtype=mx.float32) + + class FakeNestedStateCache: + @property + def state(self): + return [ + mx.array([1.0], dtype=mx.float32), + None, + {"nested": (mx.array([2.0], dtype=mx.float32),)}, + ] + + class FakeRequest: + def __init__(self): + self.req_pool_idx = None + self.mamba_pool_idx = None + self.inflight_middle_chunks = 0 + self.kv_committed_len = 0 + + class FakeTpWorker: + def __init__(self, next_token_ids): + self.result = GenerationBatchResult(next_token_ids=next_token_ids) + self.calls = [] + + def finalize_mlx_result(self, *args): + self.calls.append(args) + return self.result + + class FakeOverlapScheduler(SchedulerMlxOverlapMixin): + def __init__(self, next_token_ids): + self.tp_worker = FakeTpWorker(next_token_ids) + self.last_batch = None + self.processed_batch = None + self.processed_result = None + + def process_batch_result(self, batch, result): + self.processed_batch = batch + self.processed_result = result + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index f3a3d5f2e..174f62b14 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -41,8 +41,11 @@ from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllo from sglang.srt.mem_cache.unified_cache_components.tree_component import ( CacheTransferPhase, ComponentType, + EvictLayer, + TreeComponent, ) from sglang.srt.mem_cache.unified_radix_cache import ( + COMPONENT_REGISTRY, UnifiedRadixCache, UnifiedTreeNode, ) @@ -119,6 +122,52 @@ class CacheConfig: return "_".join(parts) +class _FakeFullComponent(TreeComponent): + component_type = ComponentType.FULL + + def create_match_validator(self, match_device_only: bool = False): + return lambda node: True + + def redistribute_on_node_split(self, new_parent, child): + return None + + def evict_component( + self, node, target: EvictLayer = EvictLayer.DEVICE + ) -> tuple[int, int]: + return 0, 0 + + def drive_eviction(self, params: EvictParams, tracker: dict[ComponentType, int]): + return None + + def acquire_component_lock(self, node, result): + return result + + def release_component_lock(self, node, params): + return None + + +class TestUnifiedRadixComponentRegistryOverride(CustomTestCase): + def test_component_registry_override_is_instance_local(self): + params = CacheInitParams( + req_to_token_pool=ReqToTokenPool( + size=2, + max_context_len=8, + device="cpu", + enable_memory_saver=False, + ), + token_to_kv_pool_allocator=None, + page_size=1, + disable=True, + tree_components=(ComponentType.FULL,), + component_registry_override={ComponentType.FULL: _FakeFullComponent}, + ) + + tree = UnifiedRadixCache(params=params) + + self.assertIsInstance(tree.components[ComponentType.FULL], _FakeFullComponent) + self.assertIsNot(COMPONENT_REGISTRY[ComponentType.FULL], _FakeFullComponent) + + def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False): """Create (tree, allocator, req_to_token_pool) from a CacheConfig.""" server_args = ServerArgs(model_path="dummy", page_size=cfg.page_size)