Create vllm_engine.py

Author: rwilliamspbg-ops

prototype/vllm_engine.py

@@ -0,0 +1,92 @@
+# prototype/vllm_engine.py (NEW)
+from typing import List, Optional
+import torch
+from transformers import AutoTokenizer
+
+class VllmInferenceEngine:
+    """Production inference engine using vLLM for optimal GPU utilization."""
+    
+    def __init__(
+        self,
+        model_id: str,
+        gpu_memory_fraction: float = 0.8,
+        max_num_seqs: int = 64,
+        tensor_parallel_size: int = 1,
+        enforce_eager: bool = False
+    ):
+        self.model_id = model_id
+        self.gpu_memory_fraction = gpu_memory_fraction
+        self.max_num_seqs = max_num_seqs
+        self.tensor_parallel_size = tensor_parallel_size
+        
+        import vllm
+        
+        # Initialize vLLM engine with optimal settings
+        self.engine = vllm.LLM(
+            model=model_id,
+            gpu_memory_fraction=self.gpu_memory_fraction,
+            max_num_seqs=self.max_num_seqs,
+            tensor_parallel_size=self.tensor_parallel_size,
+            enforce_eager=enforce_eager,  # For debugging/control
+            enable_prefix_caching=True,  # Performance optimization
+        )
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+    
+    def generate(
+        self,
+        prompts: List[str],
+        max_tokens: int = 256,
+        temperature: float = 1.0,
+        top_p: float = 0.9,
+        stop: Optional[List[str]] = None
+    ) -> List[str]:
+        """Batched generation with vLLM."""
+        
+        # Tokenize prompts
+        input_ids = self.tokenizer(
+            prompts,
+            return_tensors="pt",
+            padding=True
+        ).to(self.engine.device)
+        
+        # Generate using vLLM's optimized path
+        outputs = self.engine.generate(
+            **input_ids,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stop=stop
+        )
+        
+        # Decode and return
+        decoded = self.tokenizer.batch_decode(outputs)
+        return decoded
+    
+    def tokenized_generate(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_tokens: int = 256
+    ) -> List[str]:
+        """Generate from pre-tokenized input (for pipeline parallelism)."""
+        
+        outputs = self.engine.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_tokens=max_tokens
+        )
+        
+        decoded = self.tokenizer.batch_decode(outputs)
+        return decoded
+    
+    def get_model_profile(self) -> dict:
+        """Get vLLM engine profile for monitoring."""
+        profile = {
+            "num_active_requests": len(self.engine.get_cache_config()),
+            "cache_hit_rate": self.engine.get_cache_config().cache_hit_rate if hasattr(self.engine, 'get_cache_config') else None,
+            "gpu_memory_utilization": self.engine.get_gpu_memory_utilization(),
+            "num_tokens_seen": self.engine.get_num_prompt_tokens_seen() + self.engine.get_num_generation_tokens_seen(),
+        }
+        
+        return profile