Update torch_vgpu.py

torch_vgpu.py

@@ -35,17 +35,35 @@ class VGPUDevice:
         with vgpu.mode():
             tensor = torch.randn(2, 3)  # Will be on vGPU
     """
+    _VGPU_INSTANCES = {}  # Class-level dict to track instances
+    
     def __init__(self, vram: Optional[VirtualVRAM] = None):
         self.vram = vram or VirtualVRAM()
         self.tensor_cores = None  # Will be initialized when needed
-        self.device_name = "privateuseone"  # Use privateuseone as base device type
+        self.internal_name = "privateuseone"  # PyTorch backend name
+        self.device_name = "vgpu"  # User-facing device name
         self._register_device()
 
     def _register_device(self):
         """Register vGPU device using PyTorch's device system"""
         try:
-            # Create device instance using privateuseone backend
-            self._device = torch.device(f"{self.device_name}:0")
+            # Create internal device using privateuseone backend
+            self._device = torch.device(f"{self.internal_name}:0")
+            
+            # Store this instance for device mapping
+            VGPUDevice._VGPU_INSTANCES[self.device_name] = self
+            
+            # Register custom dispatcher for device mapping
+            def device_mapper(device_str):
+                if device_str.startswith(self.device_name):
+                    # Map vgpu -> privateuseone
+                    idx = device_str.split(":", 1)[1] if ":" in device_str else "0"
+                    return torch.device(f"{self.internal_name}:{idx}")
+                return None
+            
+            # Register the mapper with PyTorch
+            if not hasattr(torch, '_vgpu_device_mapper'):
+                torch._vgpu_device_mapper = device_mapper
             
             # Define custom operations for the device
             class VGPUAllocator:
@@ -53,10 +71,12 @@ class VGPUDevice:
                     self.vram = vram
                 
                 def __call__(self, size, dtype=None, device=None):
-                    # Create tensor on CPU first
-                    cpu_tensor = torch.empty(size, dtype=dtype, device='cpu')
-                    # Move to vGPU storage
-                    return to_vgpu(cpu_tensor, self.vram)
+                    if device is None or str(device).startswith("vgpu"):
+                        # Create tensor on CPU first
+                        cpu_tensor = torch.empty(size, dtype=dtype, device='cpu')
+                        # Move to vGPU storage
+                        return to_vgpu(cpu_tensor, self.vram)
+                    return torch.empty(size, dtype=dtype, device=device)
             
             # Set this device as the default for tensor allocation
             self._allocator = VGPUAllocator(self.vram)
@@ -66,17 +86,17 @@ class VGPUDevice:
     
     @property
     def type(self):
-        return "vgpu"  # User-facing device type name
+        return self.device_name
         
     def __str__(self):
-        return "vgpu"
+        return f"{self.device_name}:0"
         
     def __repr__(self):
-        return "vgpu"
+        return f"{self.device_name}:0"
         
     def device(self):
-        """Get the PyTorch device object"""
-        return self._device
+        """Get the PyTorch device object that maps to our vGPU"""
+        return torch.device(str(self))
     
     def mode(self):
         """Get a context manager for vGPU operations"""
@@ -121,11 +141,19 @@ def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.T
     if not isinstance(tensor, torch.Tensor):
         tensor = torch.tensor(tensor)
     
-    device = VGPUDevice(vram)
+    # Get or create vGPU device
+    if not VGPUDevice._VGPU_INSTANCES:
+        device = VGPUDevice(vram)
+    else:
+        device = next(iter(VGPUDevice._VGPU_INSTANCES.values()))
+        if vram is not None:
+            device.vram = vram
+    
+    # Move data to vRAM
     tensor_id = device._to_vram(tensor)
     result = device._from_vram(tensor_id)
     result.requires_grad = tensor.requires_grad
     
-    # Set the device correctly
-    result.data = result.data.to(device.device())
+    # Set the device using the user-facing name
+    result.data = result.data.to(f"{device.device_name}:0")
     return result