[GAUDISW-244575] Reapply Adapt OnlineDefragmenter and CacheSwapUtils for t.compile

Because of the double entrypoint of CacheSwapUtils (forward and swap functions)
torch.compile would process module and forward function while swap's self
would refer to unwrapped module. That results in the function not being run as compiled
Changes made in this patch:

 - Hide CacheSwapUtils entirely in OnlineDefragmenter. Let it be responsible for calling the module correctly
 - Moved warmup_defragmenter to defragmenter itself
 - Removed initialize function of OnlineDefragmenter, fully initialize object in init
 - Adapted unit tests for new implementations

Signed-off-by: Jan Wieczorek <jwieczorek@habana.ai>

Author: jwieczorekhabana

tests/unit_tests/test_defragmentation.py

@@ -36,7 +36,8 @@ def mock_debug_logger():
 @pytest.fixture
 def defragmenter(mock_config, mock_debug_logger):
     """Create OnlineDefragmenter instance"""
-    return OnlineDefragmenter()
+    kv_caches = ((torch.empty(0, device='meta'), torch.empty(0, device='meta')), )
+    return OnlineDefragmenter(kv_caches, block_size=0)
 
 
 class TestOnlineDefragmenter:
@@ -174,7 +175,9 @@ def test_free_blocks_generator(self, defragmenter):
     def test_defragment_disabled(self, mock_config, mock_debug_logger):
         """Test defragmentation when disabled"""
         mock_config.defrag = False
-        defrag = OnlineDefragmenter()
+
+        kv_caches = ((torch.empty(0, device='meta'), torch.empty(0, device='meta')), )
+        defrag = OnlineDefragmenter(kv_caches, 0)
 
         defrag.use_block(100)
         defrag.defragment()
@@ -198,11 +201,11 @@ def test_defragment_below_threshold(self, defragmenter):
 
         max_before = max(defragmenter.used_blocks.keys())
         defragmenter._extend_mapping_table(max_before)
-        defragmenter.cache_utils = MagicMock()
+        defragmenter._swap = MagicMock()
         defragmenter.defragment()
 
         # Should not trigger defragmentation
-        defragmenter.cache_utils.swap.assert_not_called()
+        defragmenter._swap.assert_not_called()
         assert max(defragmenter.used_blocks.keys()) == max_before
 
     def test_defragment_triggers(self, defragmenter):
@@ -214,13 +217,13 @@ def test_defragment_triggers(self, defragmenter):
             defragmenter.use_block(i)
 
         defragmenter._extend_mapping_table(102)
-        defragmenter.cache_utils = MagicMock()
+        defragmenter._swap = MagicMock()
 
         defragmenter.defragment()
 
         # Should call swap with high blocks moved to low positions
-        defragmenter.cache_utils.swap.assert_called_once()
-        args = defragmenter.cache_utils.swap.call_args[0]
+        defragmenter._swap.assert_called_once()
+        args = defragmenter._swap.call_args[0]
         to_swap = args[0]
         threshold = args[1]
 
@@ -240,62 +243,20 @@ def test_defragment_early_exit(self, defragmenter):
         defragmenter.use_block(100)
 
         defragmenter._extend_mapping_table(100)
-        defragmenter.cache_utils = MagicMock()
+        defragmenter._swap = MagicMock()
 
         defragmenter.defragment()
 
         # Free blocks: 1, 3, 4, 5...
         # Used blocks (descending): 100, 2
         # Pair (100, 1): valid swap
         # Pair (2, 3): 3 > 2, so break
-        args = defragmenter.cache_utils.swap.call_args[0]
+        args = defragmenter._swap.call_args[0]
         to_swap = args[0]
 
         assert len(to_swap) == 1
         assert to_swap[0] == (100, 1)
 
-
-class TestCacheSwapUtils:
-    """Test suite for CacheSwapUtils"""
-
-    @pytest.fixture
-    def mock_kv_caches(self):
-        """Create mock KV cache tensors"""
-        num_blocks = 100
-        block_size = 16
-        num_heads = 8
-        head_dim = 64
-        num_layers = 2
-
-        kv_caches = []
-        for _ in range(num_layers):
-            k_cache = torch.randn(num_blocks * block_size, num_heads, head_dim)
-            v_cache = torch.randn(num_blocks * block_size, num_heads, head_dim)
-            kv_caches.append((k_cache, v_cache))
-        return tuple(kv_caches)
-
-    @pytest.fixture
-    def swap_utils(self, mock_kv_caches):
-        """Create CacheSwapUtils instance"""
-        with patch('vllm_gaudi.extension.defragmentation.htorch'):
-            return CacheSwapUtils(mock_kv_caches, block_size=16)
-
-    def test_cache_swap_utils_init(self, swap_utils):
-        """Test CacheSwapUtils initialization"""
-        assert swap_utils.block_size == 16
-        assert len(swap_utils.kv_caches) == 2
-        assert swap_utils.block_slots.shape == (16, )
-        assert swap_utils.is_mla is False
-
-    def test_cache_swap_utils_mla_detection(self):
-        """Test MLA (multi-layer attention) detection"""
-        # Create MLA-style caches (no value cache)
-        mla_caches = [(torch.randn(100, 8, 64), None), (torch.randn(100, 8, 64), None)]
-
-        with patch('vllm_gaudi.extension.defragmentation.htorch'):
-            utils = CacheSwapUtils(tuple(mla_caches), block_size=16)
-            assert utils.is_mla is True
-
     def test_swap_execution(self):
         """Test swap method execution flow on HPU"""
         import habana_frameworks.torch as htorch
@@ -305,14 +266,15 @@ def test_swap_execution(self):
         num_heads = 8
         head_dim = 64
         num_layers = 2
+        DEVICE = 'hpu'
 
         kv_caches = []
         for _ in range(num_layers):
-            k_cache = torch.randn(num_blocks * block_size, num_heads, head_dim, device='hpu')
-            v_cache = torch.randn(num_blocks * block_size, num_heads, head_dim, device='hpu')
+            k_cache = torch.randn(num_blocks * block_size, num_heads, head_dim, device=DEVICE)
+            v_cache = torch.randn(num_blocks * block_size, num_heads, head_dim, device=DEVICE)
             kv_caches.append((k_cache, v_cache))
 
-        swap_utils = CacheSwapUtils(tuple(kv_caches), block_size=16)
+        defragmenter = OnlineDefragmenter(kv_caches, block_size=block_size)
 
         to_swap = [(10, 5), (20, 6)]
         threshold = 8
@@ -321,7 +283,7 @@ def test_swap_execution(self):
         orig_k_10 = kv_caches[0][0][10 * block_size:(10 + 1) * block_size].clone()
         orig_k_5 = kv_caches[0][0][5 * block_size:(5 + 1) * block_size].clone()
 
-        swap_utils.swap(to_swap, threshold)
+        defragmenter._swap(to_swap, threshold)
         htorch.core.mark_step()
 
         # Verify blocks were swapped
@@ -335,35 +297,69 @@ def test_swap_execution(self):
     def test_swap_mla_single_call(self, mock_htorch):
         """Test MLA swap only calls forward once (no value cache)"""
         mla_caches = [(torch.randn(100, 8, 64), None), (torch.randn(100, 8, 64), None)]
-        utils = CacheSwapUtils(tuple(mla_caches), block_size=16)
+        defragmenter = OnlineDefragmenter(mla_caches, block_size=16)
 
         to_swap = [(10, 5)]
         threshold = 8
 
-        with patch.object(utils, 'forward') as mock_forward:
-            utils.swap(to_swap, threshold)
+        with patch.object(defragmenter.cache_utils, 'forward') as mock_forward:
+            defragmenter._swap(to_swap, threshold)
 
             # Should only be called once for keys (no values)
             assert mock_forward.call_count == 1
 
 
+class TestCacheSwapUtils:
+    """Test suite for CacheSwapUtils"""
+
+    @pytest.fixture
+    def mock_kv_caches(self):
+        """Create mock KV cache tensors"""
+        num_blocks = 100
+        block_size = 16
+        num_heads = 8
+        head_dim = 64
+        num_layers = 2
+
+        kv_caches = []
+        for _ in range(num_layers):
+            k_cache = torch.randn(num_blocks * block_size, num_heads, head_dim)
+            v_cache = torch.randn(num_blocks * block_size, num_heads, head_dim)
+            kv_caches.append((k_cache, v_cache))
+        return tuple(kv_caches)
+
+    @pytest.fixture
+    def swap_utils(self, mock_kv_caches):
+        """Create CacheSwapUtils instance"""
+        with patch('vllm_gaudi.extension.defragmentation.htorch'):
+            return CacheSwapUtils(16, 'hpu')
+
+
 class TestDefragmentationIntegration:
     """Integration tests for defragmentation workflow"""
 
     @pytest.fixture
     def setup_defragmenter(self, mock_config, mock_debug_logger):
         """Setup defragmenter with mock caches"""
-        defrag = OnlineDefragmenter()
 
         # Create simple mock caches
         kv_caches = [(torch.zeros(1600, 8, 64), torch.zeros(1600, 8, 64)),
                      (torch.zeros(1600, 8, 64), torch.zeros(1600, 8, 64))]
 
         with patch('vllm_gaudi.extension.defragmentation.htorch'):
-            defrag.initialize(tuple(kv_caches), block_size=16)
+            defrag = OnlineDefragmenter(tuple(kv_caches), block_size=16)
 
         return defrag
 
+    def test_cache_swap_utils_mla_detection(self):
+        """Test MLA (multi-layer attention) detection"""
+        # Create MLA-style caches (no value cache)
+        mla_caches = [(torch.randn(100, 8, 64), None), (torch.randn(100, 8, 64), None)]
+
+        with patch('vllm_gaudi.extension.defragmentation.htorch'):
+            utils = OnlineDefragmenter(tuple(mla_caches), block_size=16)
+            assert utils.is_mla is True
+
     def test_full_lifecycle(self, setup_defragmenter):
         """Test complete request lifecycle with defragmentation"""
         defrag = setup_defragmenter
@@ -385,7 +381,7 @@ def test_full_lifecycle(self, setup_defragmenter):
         defrag.update_state({'req_3': [100, 101, 102]}, [])
 
         # Trigger defragmentation
-        with patch.object(defrag.cache_utils, 'swap'):
+        with patch.object(defrag, '_swap'):
             defrag.defragment()
 
     def test_mapping_persistence(self, setup_defragmenter):
@@ -398,7 +394,7 @@ def test_mapping_persistence(self, setup_defragmenter):
 
         defrag._extend_mapping_table(100)
 
-        with patch.object(defrag.cache_utils, 'swap'):
+        with patch.object(defrag, '_swap'):
             defrag.defragment()
 
         # Verify mappings exist