perf(disaggregation): cache static buffer metadata

test/registered/unit/mem_cache/test_kv_pool_buf_infos.py

@@ -0,0 +1,109 @@
+"""Unit tests for static disagg KV buffer metadata caches."""
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from sglang.srt.mem_cache.memory_pool import (
+    MHATokenToKVPool,
+    MHATokenToKVPoolFP4,
+    MLATokenToKVPool,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=6, suite="stage-a-test-cpu")
+
+
+class _BrokenBuffer:
+    def data_ptr(self):
+        raise AssertionError("buffer infos were recomputed")
+
+    @property
+    def nbytes(self):
+        raise AssertionError("buffer infos were recomputed")
+
+    def __getitem__(self, _key):
+        raise AssertionError("buffer infos were recomputed")
+
+
+class TestKVPoolBufInfos(CustomTestCase):
+    def test_mha_contiguous_buf_infos_are_cached_but_return_copies(self):
+        pool = MHATokenToKVPool(
+            size=16,
+            page_size=1,
+            dtype=torch.bfloat16,
+            head_num=2,
+            head_dim=8,
+            layer_num=2,
+            device="cpu",
+            enable_memory_saver=False,
+        )
+
+        first = pool.get_contiguous_buf_infos()
+        first[0][0] = -1
+
+        with (
+            patch.object(
+                pool,
+                "_get_key_buffer",
+                side_effect=AssertionError("key buffer infos were recomputed"),
+            ),
+            patch.object(
+                pool,
+                "_get_value_buffer",
+                side_effect=AssertionError("value buffer infos were recomputed"),
+            ),
+        ):
+            second = pool.get_contiguous_buf_infos()
+
+        self.assertNotEqual(second[0][0], -1)
+        self.assertIsNot(first[0], second[0])
+        self.assertIsNot(first[1], second[1])
+        self.assertIsNot(first[2], second[2])
+
+    def test_mla_contiguous_buf_infos_are_cached_but_return_copies(self):
+        pool = MLATokenToKVPool(
+            size=16,
+            page_size=1,
+            dtype=torch.bfloat16,
+            kv_lora_rank=8,
+            qk_rope_head_dim=4,
+            layer_num=2,
+            device="cpu",
+            enable_memory_saver=False,
+        )
+
+        first = pool.get_contiguous_buf_infos()
+        first[0][0] = -1
+        pool.kv_buffer = [_BrokenBuffer() for _ in range(pool.layer_num)]
+
+        second = pool.get_contiguous_buf_infos()
+
+        self.assertNotEqual(second[0][0], -1)
+        self.assertIsNot(first[0], second[0])
+        self.assertIsNot(first[1], second[1])
+        self.assertIsNot(first[2], second[2])
+
+    def test_mha_fp4_contiguous_buf_infos_do_not_cache_temporary_views(self):
+        pool = object.__new__(MHATokenToKVPoolFP4)
+        pool.start_layer = 0
+        pool.layer_num = 1
+        pool.page_size = 1
+
+        key_buf = torch.zeros((2, 2), dtype=torch.bfloat16)
+        value_buf = torch.zeros((2, 2), dtype=torch.bfloat16)
+        pool._get_key_buffer = MagicMock(return_value=key_buf)
+        pool._get_value_buffer = MagicMock(return_value=value_buf)
+
+        first = pool.get_contiguous_buf_infos()
+        second = pool.get_contiguous_buf_infos()
+
+        self.assertEqual(pool._get_key_buffer.call_count, 6)
+        self.assertEqual(pool._get_value_buffer.call_count, 6)
+        self.assertEqual(first, second)
+
+
+if __name__ == "__main__":
+    unittest.main()