diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index f7c5555dd43e..ba5818284e96 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -1867,6 +1867,7 @@ def get_low_ratio_index_k_dequant( self, layer_id: int, slots: Optional[torch.Tensor] = None ) -> torch.Tensor: """Index-K rows at `slots` from the layer's latent source.""" + self.wait_layer_transfer(layer_id) compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] return self._indexer_pool(compress_ratio).get_index_k_dequant( compress_layer_id, slots @@ -1877,6 +1878,7 @@ def get_low_ratio_index_k_fp4( ) -> Tuple[torch.Tensor, torch.Tensor]: """Packed fp4 index-K rows at `slots`: (payload int8 [n, 64], ue8m0 scales packed int32 [n]), the input layout of quantize_fp4_indexer_tensor.""" + self.wait_layer_transfer(layer_id) compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] return self._indexer_pool(compress_ratio).get_index_k_fp4( compress_layer_id, slots diff --git a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py b/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py index 06e677a28778..a83bb4baac17 100644 --- a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py +++ b/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py @@ -205,12 +205,18 @@ def test_indexer_access_uses_layer_ratio_and_waits_only_for_reads(self): with patch.object(pool, "wait_layer_transfer") as wait: trace.attach_mock(wait, "wait") pool.get_index_k_fp4_payload_buffer(2) + pool.get_low_ratio_index_k_dequant(2, "slots") + pool.get_low_ratio_index_k_fp4(2, "slots") pool.set_index_k_fp4(2, "loc", "cache") self.assertEqual( trace.mock_calls, [ unittest.mock.call.wait(2), unittest.mock.call.indexer.get_index_k_fp4_payload_buffer(1), + unittest.mock.call.wait(2), + unittest.mock.call.indexer.get_index_k_dequant(1, "slots"), + unittest.mock.call.wait(2), + unittest.mock.call.indexer.get_index_k_fp4(1, "slots"), unittest.mock.call.indexer.set_index_fp4(1, "loc", "cache"), ], )