Upload 2 files

test_ai_integration_http.py

@@ -1,203 +1,208 @@
-"""
-Test Llama-2-7b-instruct model integration with vGPU.
-Configure PyTorch to use vGPU as device for text generation.
-"""
-import logging
-import os
-import time
-from contextlib import contextmanager
-from typing import Any, Optional
-
-import torch
-from transformers import pipeline
-from virtual_vram import VirtualVRAM
-from http_storage import HTTPGPUStorage
-from torch_vgpu import VGPUDevice, to_vgpu
-
-def setup_vgpu():
-    """Setup vGPU device"""
-    try:
-        # Create and register vGPU device
-        vgpu = VGPUDevice()
-        device = vgpu.device()
-        
-        # Set as default device for tensor operations
-        torch.set_default_device(device)
-        
-        return device
-
-    except Exception as e:
-        logging.error(f"vGPU setup failed: {str(e)}")
-        raise
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
-
-@contextmanager
-def gpu_context():
-    """Context manager for vGPU resources"""
-    storage = None
-    try:
-        storage = HTTPGPUStorage()
-        yield storage
-    finally:
-        if storage:
-            storage.close()
-            logger.info("vGPU resources cleaned up")
-
-def get_model_size(model):
-    """Calculate model size in parameters and memory footprint"""
-    param_size = 0
-    for param in model.parameters():
-        param_size += param.nelement() * param.element_size()
-    buffer_size = 0
-    for buffer in model.buffers():
-        buffer_size += buffer.nelement() * buffer.element_size()
-    return param_size + buffer_size
-
-def prepare_prompt(instruction: str) -> str:
-    """Prepare a prompt for Llama-2 using its chat format."""
-    # Format: <s>[INST] instruction [/INST] assistant response </s>[INST] ...
-    return f"<s>[INST] {instruction} [/INST]"
-
-def test_ai_integration_http():
-    """Test GPT OSS model on vGPU with text generation"""
-    logger.info("Starting vGPU text generation test")
-    
-    status = {
-        'pipeline_loaded': False,
-        'model_on_vgpu': False,
-        'generation_complete': False,
-        'cleanup_success': False
-    }
-    
-    with gpu_context() as storage:
-        try:
-            # Initialize vRAM with monitoring
-            initial_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
-            vram = VirtualVRAM(size_gb=None, storage=storage)
-            
-            # Initialize vGPU device
-            device = setup_vgpu()
-            logger.info(f"vGPU initialized with device {device}")
-            
-            # Load model using pipeline
-            model_id = "openai/gpt-oss-20b"
-            logger.info(f"Loading {model_id}")
-            
-            try:
-                # Disable transformers logging temporarily
-                transformers_logger = logging.getLogger("transformers")
-                original_level = transformers_logger.level
-                transformers_logger.setLevel(logging.ERROR)
-                
-                try:
-                    # Create pipeline
-                    # Create pipeline with vGPU device
-                    pipe = pipeline(
-                        "text-generation",
-                        model=model_id,
-                        torch_dtype=torch.float32,  # Use full precision
-                        device=device,  # Use our vGPU device
-                        use_safetensors=True,
-                        trust_remote_code=True
-                    )
-                    status['pipeline_loaded'] = True
-                    
-                    # Move pipeline model to vGPU
-                    pipe.model = to_vgpu(pipe.model, vram=vram)
-                    pipe.model.eval()
-                    status['model_on_vgpu'] = True
-                    
-                    # Log model details
-                    logger.info(f"Pipeline created with model: {model_id}")
-                    
-                    # Log model size
-                    model_size = get_model_size(pipe.model)
-                    logger.info(f"Model loaded: {model_size/1e9:.2f} GB in parameters")
-                    logger.info(f"Model architecture: {pipe.model.__class__.__name__}")
-                    
-                    # Verify model location
-                    with torch.device(device):
-                        current_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
-                        logger.info(f"Model memory usage: {(current_mem - initial_mem)/1e9:.2f} GB")
-                
-                finally:
-                    # Restore original logging level
-                    transformers_logger.setLevel(original_level)
-                    
-            except Exception as e:
-                logger.error(f"Model loading failed: {str(e)}")
-                raise
-            except Exception as e:
-                logger.error(f"Model transfer to vGPU failed: {str(e)}")
-                raise
-            
-            # Run text generation
-            logger.info("Running text generation...")
-            start = time.time()
-            peak_mem = initial_mem
-            
-            try:
-                # Prepare input prompt
-                prompt = "Explain how virtual GPUs work in simple terms."
-                
-                with torch.no_grad():
-                    # Generate text
-                    outputs = pipe(
-                        prompt,
-                        max_new_tokens=256,
-                        temperature=0.7,
-                        top_p=0.95,
-                        top_k=40,
-                        num_beams=1,
-                        do_sample=True,
-                        return_full_text=True
-                    )
-                    
-                    if hasattr(storage, 'get_used_memory'):
-                        peak_mem = max(peak_mem, storage.get_used_memory())
-                    
-                    inference_time = time.time() - start
-                    status['generation_complete'] = True
-                    
-                    # Log performance metrics
-                    logger.info(f"\nGeneration stats:")
-                    logger.info(f"- Time: {inference_time:.4f}s")
-                    logger.info(f"- Memory peak: {(peak_mem - initial_mem)/1e9:.2f} GB")
-                    logger.info(f"- Generated text: {outputs[0]['generated_text']}")
-                        
-            except Exception as e:
-                logger.error(f"Text generation failed: {str(e)}")
-                raise
-            
-        except Exception as e:
-            logger.error(f"Test failed: {str(e)}")
-            raise
-        finally:
-            # Cleanup and status report
-            try:
-                if 'pipe' in locals():
-                    del pipe
-                if 'outputs' in locals():
-                    del outputs
-                torch.cuda.empty_cache() if hasattr(torch, 'cuda') else None
-                status['cleanup_success'] = True
-            except Exception as e:
-                logger.error(f"Cleanup error: {str(e)}")
-            
-            logger.info("\nTest Summary:")
-            for key, value in status.items():
-                logger.info(f"- {key}: {'✓' if value else '✗'}")
-            
-            final_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
-            if final_mem > initial_mem:
-                logger.warning(f"Memory leak detected: {(final_mem - initial_mem)/1e6:.2f} MB")
-            
-if __name__ == "__main__":
-    test_ai_integration_http()
-
+"""
+Test Llama-2-7b-instruct model integration with vGPU.
+Configure PyTorch to use vGPU as device for text generation.
+"""
+import logging
+import os
+import time
+from contextlib import contextmanager
+from typing import Any, Optional
+
+import torch
+from transformers import pipeline
+from virtual_vram import VirtualVRAM
+from http_storage import HTTPGPUStorage
+from torch_vgpu import VGPUDevice, to_vgpu
+
+def setup_vgpu():
+    """Setup vGPU device"""
+    try:
+        # Initialize the backend first
+        from torch_vgpu import init_vgpu_backend, VGPUDevice
+        if not init_vgpu_backend():
+            raise RuntimeError("Failed to initialize vGPU backend")
+        
+        # Create and register vGPU device
+        vgpu = VGPUDevice()
+        device = vgpu.device()
+        
+        # Set as default device for tensor operations
+        torch.set_default_device(device)
+        
+        return device
+
+    except Exception as e:
+        logging.error(f"vGPU setup failed: {str(e)}")
+        raise
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+@contextmanager
+def gpu_context():
+    """Context manager for vGPU resources"""
+    storage = None
+    try:
+        storage = HTTPGPUStorage()
+        yield storage
+    finally:
+        if storage:
+            storage.close()
+            logger.info("vGPU resources cleaned up")
+
+def get_model_size(model):
+    """Calculate model size in parameters and memory footprint"""
+    param_size = 0
+    for param in model.parameters():
+        param_size += param.nelement() * param.element_size()
+    buffer_size = 0
+    for buffer in model.buffers():
+        buffer_size += buffer.nelement() * buffer.element_size()
+    return param_size + buffer_size
+
+def prepare_prompt(instruction: str) -> str:
+    """Prepare a prompt for Llama-2 using its chat format."""
+    # Format: <s>[INST] instruction [/INST] assistant response </s>[INST] ...
+    return f"<s>[INST] {instruction} [/INST]"
+
+def test_ai_integration_http():
+    """Test GPT OSS model on vGPU with text generation"""
+    logger.info("Starting vGPU text generation test")
+    
+    status = {
+        'pipeline_loaded': False,
+        'model_on_vgpu': False,
+        'generation_complete': False,
+        'cleanup_success': False
+    }
+    
+    with gpu_context() as storage:
+        try:
+            # Initialize vRAM with monitoring
+            initial_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
+            vram = VirtualVRAM(size_gb=None, storage=storage)
+            
+            # Initialize vGPU device
+            device = setup_vgpu()
+            logger.info(f"vGPU initialized with device {device}")
+            
+            # Load model using pipeline
+            model_id = "openai/gpt-oss-20b"
+            logger.info(f"Loading {model_id}")
+            
+            try:
+                # Disable transformers logging temporarily
+                transformers_logger = logging.getLogger("transformers")
+                original_level = transformers_logger.level
+                transformers_logger.setLevel(logging.ERROR)
+                
+                try:
+                    # Create pipeline
+                    # Create pipeline with vGPU device
+                    pipe = pipeline(
+                        "text-generation",
+                        model=model_id,
+                        torch_dtype=torch.float32,  # Use full precision
+                        device=device,  # Use our vGPU device
+                        use_safetensors=True,
+                        trust_remote_code=True
+                    )
+                    status['pipeline_loaded'] = True
+                    
+                    # Move pipeline model to vGPU
+                    pipe.model = to_vgpu(pipe.model, vram=vram)
+                    pipe.model.eval()
+                    status['model_on_vgpu'] = True
+                    
+                    # Log model details
+                    logger.info(f"Pipeline created with model: {model_id}")
+                    
+                    # Log model size
+                    model_size = get_model_size(pipe.model)
+                    logger.info(f"Model loaded: {model_size/1e9:.2f} GB in parameters")
+                    logger.info(f"Model architecture: {pipe.model.__class__.__name__}")
+                    
+                    # Verify model location
+                    with torch.device(device):
+                        current_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
+                        logger.info(f"Model memory usage: {(current_mem - initial_mem)/1e9:.2f} GB")
+                
+                finally:
+                    # Restore original logging level
+                    transformers_logger.setLevel(original_level)
+                    
+            except Exception as e:
+                logger.error(f"Model loading failed: {str(e)}")
+                raise
+            except Exception as e:
+                logger.error(f"Model transfer to vGPU failed: {str(e)}")
+                raise
+            
+            # Run text generation
+            logger.info("Running text generation...")
+            start = time.time()
+            peak_mem = initial_mem
+            
+            try:
+                # Prepare input prompt
+                prompt = "Explain how virtual GPUs work in simple terms."
+                
+                with torch.no_grad():
+                    # Generate text
+                    outputs = pipe(
+                        prompt,
+                        max_new_tokens=256,
+                        temperature=0.7,
+                        top_p=0.95,
+                        top_k=40,
+                        num_beams=1,
+                        do_sample=True,
+                        return_full_text=True
+                    )
+                    
+                    if hasattr(storage, 'get_used_memory'):
+                        peak_mem = max(peak_mem, storage.get_used_memory())
+                    
+                    inference_time = time.time() - start
+                    status['generation_complete'] = True
+                    
+                    # Log performance metrics
+                    logger.info(f"\nGeneration stats:")
+                    logger.info(f"- Time: {inference_time:.4f}s")
+                    logger.info(f"- Memory peak: {(peak_mem - initial_mem)/1e9:.2f} GB")
+                    logger.info(f"- Generated text: {outputs[0]['generated_text']}")
+                        
+            except Exception as e:
+                logger.error(f"Text generation failed: {str(e)}")
+                raise
+            
+        except Exception as e:
+            logger.error(f"Test failed: {str(e)}")
+            raise
+        finally:
+            # Cleanup and status report
+            try:
+                if 'pipe' in locals():
+                    del pipe
+                if 'outputs' in locals():
+                    del outputs
+                torch.cuda.empty_cache() if hasattr(torch, 'cuda') else None
+                status['cleanup_success'] = True
+            except Exception as e:
+                logger.error(f"Cleanup error: {str(e)}")
+            
+            logger.info("\nTest Summary:")
+            for key, value in status.items():
+                logger.info(f"- {key}: {'✓' if value else '✗'}")
+            
+            final_mem = storage.get_used_memory() if hasattr(storage, 'get_used_memory') else 0
+            if final_mem > initial_mem:
+                logger.warning(f"Memory leak detected: {(final_mem - initial_mem)/1e6:.2f} MB")
+            
+if __name__ == "__main__":
+    test_ai_integration_http()
+

torch_vgpu.py

@@ -1,178 +1,182 @@
-"""
-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
-
-# Initialize custom backend
-def init_vgpu_backend():
-    try:
-        # First rename the backend
-        torch.utils.rename_privateuse1_backend("vgpu")
-        
-        # Then generate all the necessary methods
-        torch.utils.generate_methods_for_privateuse1_backend(
-            for_tensor=True,
-            for_module=True,
-            for_packed_sequence=True,
-            for_storage=True
-        )
-        
-        # Register our custom library
-        lib = Library("vgpu", "DEF")
-        lib.define("custom_op(Tensor self) -> Tensor")
-        
-        @impl("vgpu", "custom_op", "Tensor")
-        def custom_op_impl(tensor):
-            return tensor.clone()
-            
-        return True
-    except Exception as e:
-        print(f"Backend initialization warning: {e}")
-        return False
-
-# Initialize the backend
-VGPU_BACKEND_INITIALIZED = init_vgpu_backend()
-
-class VGPUStorage(torch.Storage):
-    """Custom storage class that uses our virtual VRAM"""
-    
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.vram = kwargs.get('vram')
-        if not self.vram:
-            from virtual_vram import VirtualVRAM
-            self.vram = VirtualVRAM()
-        self.tensor_id = kwargs.get('tensor_id', f"tensor_{id(self)}")
-    
-    def _new_shared(self, size):
-        return VGPUStorage(size, vram=self.vram)
-
-class VGPUTensor:
-    """Tensor implementation that uses vGPU for computations"""
-    @staticmethod
-    def __new__(cls, elem):
-        return torch.Tensor._make_subclass(cls, elem, elem.requires_grad)
-
-class VGPUDevice:
-    """
-    Custom PyTorch device implementation that routes operations through vGPU.
-    Usage:
-        vgpu = 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 = "vgpu"  # Both internal and user-facing name
-        self._register_device()
-
-    def _register_device(self):
-        """Register vGPU device using PyTorch's device system"""
-        try:
-            if not VGPU_BACKEND_INITIALIZED:
-                raise RuntimeError("VGPU backend not properly initialized")
-                
-            # Create device with explicit index
-            self._device = torch.device("vgpu")
-            
-            # Store this instance for reuse
-            VGPUDevice._VGPU_INSTANCES[self.device_name] = self
-            
-            # Define custom operations for the device
-            class VGPUAllocator:
-                def __init__(self, vram, device):
-                    self.vram = vram
-                    self.device = device
-                
-                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)
-            
-            # Set up allocator
-            self._allocator = VGPUAllocator(self.vram, self._device)
-            
-        except Exception as e:
-            raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
-    
-    @property
-    def type(self):
-        return self.internal_name
-        
-    def __str__(self):
-        return f"{self.internal_name}:0"
-        
-    def __repr__(self):
-        return f"vgpu(device='{self.internal_name}:0')"
-        
-    def device(self):
-        """Get the PyTorch device object that maps to our vGPU"""
-        return self._device  # Return the already created device object
-    
-    def mode(self):
-        """Get a context manager for vGPU operations"""
-        return torch.device(self._device)
-    
-    def _init_tensor_cores(self):
-        if self.tensor_cores is None:
-            from tensor_core import TensorCoreArray
-            self.tensor_cores = TensorCoreArray()
-            
-    def _to_vram(self, tensor: torch.Tensor) -> str:
-        """Store tensor data in virtual VRAM"""
-        tensor_id = f"tensor_{id(tensor)}"
-        data = tensor.detach().cpu().numpy()
-        self.vram.storage.store_tensor(tensor_id, data)
-        return tensor_id
-        
-    def _from_vram(self, tensor_id: str) -> torch.Tensor:
-        """Retrieve tensor data from virtual VRAM"""
-        data = self.vram.storage.load_tensor(tensor_id)
-        return torch.from_numpy(data)
-
-    def matmul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        """Matrix multiplication using tensor cores"""
-        self._init_tensor_cores()
-        
-        # Store inputs in VRAM
-        a_id = self._to_vram(a)
-        b_id = self._to_vram(b)
-        
-        # Perform matmul using tensor cores
-        result = self.tensor_cores.matmul(
-            self.vram.storage.load_tensor(a_id),
-            self.vram.storage.load_tensor(b_id)
-        )
-        
-        # Create new tensor with result
-        return torch.from_numpy(result)
-
-def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
-    """Move a tensor to vGPU device"""
-    if not isinstance(tensor, torch.Tensor):
-        tensor = torch.tensor(tensor)
-    
-    # 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 using the internal name
-    result.data = result.data.to(device._device)
-    return result
+"""
+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
+
+# Global flag for backend initialization
+VGPU_BACKEND_INITIALIZED = False
+
+def init_vgpu_backend():
+    """Initialize the vGPU backend. Must be called before creating any VGPUDevice instances."""
+    global VGPU_BACKEND_INITIALIZED
+    try:
+        if not VGPU_BACKEND_INITIALIZED:
+            # First rename the backend
+            torch.utils.rename_privateuse1_backend("vgpu")
+            
+            # Then generate all the necessary methods
+            torch.utils.generate_methods_for_privateuse1_backend(
+                for_tensor=True,
+                for_module=True,
+                for_packed_sequence=True,
+                for_storage=True
+            )
+            
+            # Register our custom library
+            lib = Library("vgpu", "DEF")
+            lib.define("custom_op(Tensor self) -> Tensor")
+            
+            @impl("vgpu", "custom_op", "Tensor")
+            def custom_op_impl(tensor):
+                return tensor.clone()
+                
+            VGPU_BACKEND_INITIALIZED = True
+            
+        return VGPU_BACKEND_INITIALIZED
+    except Exception as e:
+        print(f"Backend initialization warning: {e}")
+        return False
+
+class VGPUStorage(torch.Storage):
+    """Custom storage class that uses our virtual VRAM"""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.vram = kwargs.get('vram')
+        if not self.vram:
+            from virtual_vram import VirtualVRAM
+            self.vram = VirtualVRAM()
+        self.tensor_id = kwargs.get('tensor_id', f"tensor_{id(self)}")
+    
+    def _new_shared(self, size):
+        return VGPUStorage(size, vram=self.vram)
+
+class VGPUTensor:
+    """Tensor implementation that uses vGPU for computations"""
+    @staticmethod
+    def __new__(cls, elem):
+        return torch.Tensor._make_subclass(cls, elem, elem.requires_grad)
+
+class VGPUDevice:
+    """
+    Custom PyTorch device implementation that routes operations through vGPU.
+    Usage:
+        vgpu = 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 = "vgpu"  # Both internal and user-facing name
+        self._register_device()
+
+    def _register_device(self):
+        """Register vGPU device using PyTorch's device system"""
+        try:
+            if not VGPU_BACKEND_INITIALIZED:
+                raise RuntimeError("VGPU backend not properly initialized")
+                
+            # Create device with explicit index
+            self._device = torch.device("vgpu")
+            
+            # Store this instance for reuse
+            VGPUDevice._VGPU_INSTANCES[self.device_name] = self
+            
+            # Define custom operations for the device
+            class VGPUAllocator:
+                def __init__(self, vram, device):
+                    self.vram = vram
+                    self.device = device
+                
+                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)
+            
+            # Set up allocator
+            self._allocator = VGPUAllocator(self.vram, self._device)
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
+    
+    @property
+    def type(self):
+        return self.internal_name
+        
+    def __str__(self):
+        return f"{self.internal_name}:0"
+        
+    def __repr__(self):
+        return f"vgpu(device='{self.internal_name}:0')"
+        
+    def device(self):
+        """Get the PyTorch device object that maps to our vGPU"""
+        return self._device  # Return the already created device object
+    
+    def mode(self):
+        """Get a context manager for vGPU operations"""
+        return torch.device(self._device)
+    
+    def _init_tensor_cores(self):
+        if self.tensor_cores is None:
+            from tensor_core import TensorCoreArray
+            self.tensor_cores = TensorCoreArray()
+            
+    def _to_vram(self, tensor: torch.Tensor) -> str:
+        """Store tensor data in virtual VRAM"""
+        tensor_id = f"tensor_{id(tensor)}"
+        data = tensor.detach().cpu().numpy()
+        self.vram.storage.store_tensor(tensor_id, data)
+        return tensor_id
+        
+    def _from_vram(self, tensor_id: str) -> torch.Tensor:
+        """Retrieve tensor data from virtual VRAM"""
+        data = self.vram.storage.load_tensor(tensor_id)
+        return torch.from_numpy(data)
+
+    def matmul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """Matrix multiplication using tensor cores"""
+        self._init_tensor_cores()
+        
+        # Store inputs in VRAM
+        a_id = self._to_vram(a)
+        b_id = self._to_vram(b)
+        
+        # Perform matmul using tensor cores
+        result = self.tensor_cores.matmul(
+            self.vram.storage.load_tensor(a_id),
+            self.vram.storage.load_tensor(b_id)
+        )
+        
+        # Create new tensor with result
+        return torch.from_numpy(result)
+
+def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
+    """Move a tensor to vGPU device"""
+    if not isinstance(tensor, torch.Tensor):
+        tensor = torch.tensor(tensor)
+    
+    # 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 using the internal name
+    result.data = result.data.to(device._device)
+    return result