Update torch_vgpu.py

torch_vgpu.py

@@ -2,6 +2,7 @@
 Custom PyTorch device implementation that routes operations through our virtual GPU.
 """
 import torch
+from torch.library import Library, impl
 from typing import Optional, Union, Tuple
 import numpy as np
 from virtual_vram import VirtualVRAM
@@ -41,9 +42,35 @@ class VGPUDevice:
         self._register_device()
         
     def _register_device(self):
-        """Register vGPU device using privateuse1 backend"""
+        """Register vGPU device using torch.library"""
         try:
-            torch._C._dispatch._rename_privateuse1_backend(self.device_name)
+            # Create library for vGPU backend
+            lib = torch.library.Library(self.device_name, "IMPL")
+            
+            # Register basic tensor operations
+            @torch.library.impl(lib, "aten::empty.memory_format")
+            def empty_impl(size, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None):
+                # Create empty tensor in CPU and move to vGPU
+                cpu_tensor = torch.empty(size, dtype=dtype, device='cpu')
+                return to_vgpu(cpu_tensor, self.vram)
+            
+            @torch.library.impl(lib, "aten::add.Tensor")
+            def add_impl(self, other):
+                # Custom implementation of add operation
+                # Move tensors to CPU, add, then back to vGPU
+                cpu_result = self.cpu() + other.cpu()
+                return to_vgpu(cpu_result, self.vram)
+                
+            @torch.library.impl(lib, "aten::copy_")
+            def copy_impl(self, src, non_blocking=False):
+                # Handle tensor copy operations
+                if not isinstance(src, torch.Tensor):
+                    src = torch.tensor(src)
+                return to_vgpu(src.cpu(), self.vram)
+            
+            # Get device after registration
+            self._device = torch.device(f"{self.device_name}:0")
+                
         except Exception as e:
             raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
     
@@ -59,32 +86,26 @@ class VGPUDevice:
         
     def device(self):
         """Get the PyTorch device object"""
-        return torch.device(str(self))
+        return self._device
     
     def mode(self):
         """Get a context manager for vGPU operations"""
-        from torch.utils._python_dispatch import TorchFunctionMode
-        
-        class _VGPUMode(TorchFunctionMode):
+        class _VGPUMode:
             def __init__(self, device):
                 self.device = device
             
-            def __torch_function__(self, func, types, args=(), kwargs=None):
-                kwargs = kwargs or {}
+            def __enter__(self):
+                return self
                 
-                # Handle tensor creation and device placement
-                if func is torch.tensor or 'device' in kwargs:
-                    kwargs['device'] = str(self.device)
-                
-                # Handle tensor operations
-                new_args = []
-                for arg in args:
-                    if isinstance(arg, torch.Tensor):
-                        if not str(arg.device).startswith(self.device.device_name):
-                            arg = to_vgpu(arg, self.device.vram)
-                    new_args.append(arg)
-                
-                return func(*new_args, **kwargs)
+            def __exit__(self, exc_type, exc_val, exc_tb):
+                pass
+            
+            def __call__(self, fn):
+                def wrapped(*args, **kwargs):
+                    if 'device' in kwargs:
+                        kwargs['device'] = str(self.device)
+                    return fn(*args, **kwargs)
+                return wrapped
         
         return _VGPUMode(self)
     
@@ -126,9 +147,6 @@ def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.T
     """Move a tensor to vGPU device"""
     device = VGPUDevice(vram)
     tensor_id = device._to_vram(tensor)
-    return VGPUTensor(device._from_vram(tensor_id))
-
-def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
-    """Helper function to move tensors to vGPU"""
-    device = VGPUDevice(vram)
-    return tensor.to(device=device)
+    result = device._from_vram(tensor_id)
+    result.requires_grad = tensor.requires_grad
+    return result