Update gpu_arch.py

gpu_arch.py

@@ -295,7 +295,7 @@ class GPUMemoryHierarchy:
 
 
 class Chip:
-    def __init__(self, chip_id, num_sms=1500, vram_size_gb=16, db_path="gpu_state.db"):
+    def __init__(self, chip_id, num_sms=1500, vram_size_gb=16, db_path="gpu_state.db", storage=None):
         self.chip_id = chip_id
         self.db = GPUStateDB(db_path)
         # Handle unlimited VRAM case (when vram_size_gb is None)
@@ -303,9 +303,10 @@ class Chip:
         self.gpu_mem = GPUMemoryHierarchy(num_sms=num_sms, global_mem_size_bytes=global_mem_size_bytes, chip_id=chip_id, db=self.db)
         self.sm_ids = list(range(num_sms))
         self.connected_chips = []
-        self.ai_accelerator = AIAccelerator() # Instantiate AIAccelerator
-        self.custom_vram = CustomVRAM(self.gpu_mem.global_mem) # Create CustomVRAM instance
-        self.ai_accelerator.set_vram(self.custom_vram) # Set VRAM for AIAccelerator
+        self.storage = storage  # Store shared WebSocket storage
+        self.ai_accelerator = AIAccelerator(storage=storage)  # Pass shared storage to accelerator
+        self.custom_vram = CustomVRAM(self.gpu_mem.global_mem)  # Create CustomVRAM instance
+        self.ai_accelerator.set_vram(self.custom_vram)  # Set VRAM for AIAccelerator
 
     def get_sm(self, sm_id):
         return StreamingMultiprocessor(sm_id, self.chip_id, self.db)