[HiCache] Fix sparse hybrid transfer layer IDs (#37870)

Co-authored-by: Seokhoon Kang <sh.kang@postech.ac.kr>

test/registered/unit/mem_cache/test_dsv4_hicache_l2.py

@@ -254,7 +254,7 @@ class TestDSV4PoolAssembly(CustomTestCase):
             pp_cache_group=None,
         )
         mappings = assembler._DeepSeekV4LayerMappings(
-            transfer_layer_num=1,
+            transfer_layer_id_max=1,
             full={0: 0},
             swa={},
             c4={0: 0},

test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py

@@ -238,7 +238,7 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
             start_event=object(), finish_event=object(), timing_enabled=False
         )
         controller.l2_transfer_engine.submit_host_to_device.return_value = completion
-        controller.layer_num = 2
+        controller.transfer_layer_id_max = 2
         controller.ack_load_queue = []
 
         self.assertEqual(HybridCacheController.start_loading(controller), 0)
@@ -344,7 +344,9 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
             layer_mapper={1: 0, 3: 1}.get,
         )
         with mock.patch.object(transfer_module, "device_module", _FakeDeviceModule):
-            L2TransferEngine("kernel").submit_host_to_device([transfer], layer_num=4)
+            L2TransferEngine("kernel").submit_host_to_device(
+                [transfer], transfer_layer_id_max=4
+            )
 
         self.assertEqual(
             [
@@ -369,7 +371,7 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
             anchor_entry=entry,
             entry_map={entry.name: entry},
         )
-        controller.layer_num = 2
+        controller.transfer_layer_id_max = 2
 
         self.assertEqual(
             len(controller._l2_transfers(_indices(0, 2), _indices(2, 4))), 1
@@ -380,7 +382,9 @@ class TestHiCacheStagedWriteBackDispatch(CustomTestCase):
         self.assertFalse(transfers[0].is_draft)
         self.assertTrue(transfers[1].is_draft)
         with mock.patch.object(transfer_module, "device_module", _FakeDeviceModule):
-            L2TransferEngine("kernel").submit_host_to_device(transfers, layer_num=2)
+            L2TransferEngine("kernel").submit_host_to_device(
+                transfers, transfer_layer_id_max=2
+            )
         self.assertEqual(
             [
                 call.args[3]

test/registered/unit/mem_cache/test_hybrid_pool_assembler.py

@@ -15,6 +15,7 @@ from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
     _split_hicache_size,
     _SwaStrategy,
     build_full_draft_pools,
+    build_hybrid_swa_group,
 )
 from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -153,7 +154,10 @@ class TestHybridStageLayerMappings(CustomTestCase):
                     with patch.object(
                         hybrid_pool_assembler,
                         builder_name,
-                        return_value=(MagicMock(), object()),
+                        return_value=(
+                            MagicMock(),
+                            SimpleNamespace(transfer_layer_id_max=4),
+                        ),
                     ) as build_stack:
                         result = strategy_cls().build(
                             cache=SimpleNamespace(page_size=1),
@@ -169,7 +173,7 @@ class TestHybridStageLayerMappings(CustomTestCase):
                             build_stack.call_args.kwargs[f"{name}_layer_mapping"],
                             mapping,
                         )
-                    self.assertEqual(result.transfer_layer_num, 4)
+                    self.assertEqual(result.cache_controller.transfer_layer_id_max, 4)
                     self.assertEqual(
                         kvcache.full_attention_layer_id_mapping, global_maps["full"]
                     )
@@ -229,5 +233,44 @@ class TestDraftSidecarPoolDispatch(CustomTestCase):
         self.assertIs(entries[0].host_pool, draft_host_pool)
 
 
+_ASSEMBLER = "sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler."
+
+
+class TestTransferLayerSpan(CustomTestCase):
+    """``transfer_layer_id_max`` must span global layer ids, not count the mapped ones.
+
+    A hybrid model with an uncached layer type keys its mappings non-contiguously,
+    and the per-layer transfer loop then never reaches the high layer ids.
+    """
+
+    def test_pool_entries_span_the_highest_global_layer_id(self):
+        # Global ids 0/2/4/6 with holes between them, the shape NemotronH's
+        # cache-ineligible MLP layers produce: 4 mapped layers spanning 7 ids.
+        full_layer_mapping = {0: 0, 6: 1}
+        swa_layer_mapping = {2: 0, 4: 1}
+
+        with (
+            patch(_ASSEMBLER + "build_kv_host_pool"),
+            patch(_ASSEMBLER + "HostPoolGroup"),
+            patch(_ASSEMBLER + "build_pool_entry") as build_pool_entry,
+        ):
+            build_hybrid_swa_group(
+                page_size=64,
+                full_kv_pool=MagicMock(),
+                swa_kv_pool=MagicMock(),
+                full_layer_mapping=full_layer_mapping,
+                swa_layer_mapping=swa_layer_mapping,
+                use_mla=False,
+            )
+
+        self.assertEqual(
+            [
+                c.kwargs["transfer_layer_id_max"]
+                for c in build_pool_entry.call_args_list
+            ],
+            [7, 7],
+        )
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_radix_hicache_dispatch.py

@@ -115,7 +115,6 @@ class TestUnifiedRadixHiCacheDispatch(unittest.TestCase):
         self.assertIs(result.cache_controller, cache_controller)
         self.assertIs(result.component_host_pools[FULL], kv_host_pool)
         self.assertEqual(result.pools_desc, "KV + INDEXER(k-only)")
-        self.assertEqual(result.transfer_layer_num, 8)
         self.assertEqual(len(result.sidecars), 1)
         self.assertEqual(result.sidecars[0].pool_name, PoolName.INDEXER)
         self.assertEqual(result.sidecars[0].indices_from_pool, PoolName.KV)
@@ -189,7 +188,6 @@ class TestApplyStackResult(unittest.TestCase):
             component_host_pools={FULL: full_host, SWA: swa_host, MAMBA: mamba_host},
             sidecars=[sidecar],
             register_req_to_token_counter=True,
-            transfer_layer_num=8,
             pools_desc="KV + SWA + MAMBA",
         )
 
@@ -221,7 +219,6 @@ class TestApplyStackResult(unittest.TestCase):
             component_host_pools={FULL: MagicMock()},
             sidecars=[],
             register_req_to_token_counter=False,
-            transfer_layer_num=1,
             pools_desc="KV",
         )