Fix auto merge & add back get_flat_data_by_layer (#4393)

2025-03-13 08:46:25 -07:00
parent 8e66fbecee
commit a5a892ffd3
3 changed files with 19 additions and 2 deletions
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -591,6 +591,9 @@ class MHATokenToKVPoolHost:
    def get_flat_data(self, indices):
        return self.kv_buffer[:, :, indices]

+    def get_flat_data_by_layer(self, indices, layer_id):
+        return self.kv_buffer[:, layer_id, indices]
+
    def assign_flat_data(self, indices, flat_data):
        self.kv_buffer[:, :, indices] = flat_data