sglang/test/registered/unit/hardware_backend/mlx/test_attention_patching.py

"""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()