Minor Cleanup

Author: vivekgoe

vllm_gaudi/ops/hpu_lora.py

@@ -1,7 +1,6 @@
 import torch
 import torch.nn.functional as F
 from vllm.model_executor.custom_op import CustomOp
-from vllm.lora import layers
 from vllm.lora.layers import VocabParallelEmbeddingWithLoRA
 
 

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -442,10 +442,9 @@ def generate_proposals(self, *args, **kwargs):
 
 
 def _maybe_wrap_in_hpu_graph(*args, **kwargs):
-    '''return htorch.hpu.wrap_in_hpu_graph(
+    return htorch.hpu.wrap_in_hpu_graph(
         HpuModelAdapter(*args, **kwargs), disable_tensor_cache=True
-    ) if htorch.utils.internal.is_lazy() else HpuModelAdapter(*args, **kwargs)'''
-    return  HpuModelAdapter(*args, **kwargs)
+    ) if htorch.utils.internal.is_lazy() else HpuModelAdapter(*args, **kwargs)
 
 
 def subtuple(obj: object,
@@ -619,7 +618,7 @@ def __init__(
         self.use_hpu_graph = not self.model_config.enforce_eager
         self.max_batch_size = self.scheduler_config.max_num_seqs
         self.max_num_seqs = self.scheduler_config.max_num_seqs
-        self.max_prefill_batch_size = 2  # TODO(kzawora): add knob for that
+        self.max_prefill_batch_size = 1  # TODO(kzawora): add knob for that
         self.seen_configs: set = set()
         self.max_num_batched_tokens = \
             self.scheduler_config.max_num_batched_tokens