Fix generator vllm after rebase

Author: MohammedTaherMcW

models/tt_transformers/tt/generator_vllm.py

@@ -459,16 +459,14 @@ def prefill_forward(self, *args, **kwargs):
             tokens[i][prompt_lens[i] :] = pad_token_id
         pixel_values = None
 
-        if hasattr(data[0], "pixel_values"):
-            # If inputs is a list of objects with pixel_values, concatenate them
-            pixel_values = torch.concat([im.pixel_values for im in data if hasattr(im, "pixel_values")], dim=0)
+        if any(hasattr(d, "pixel_values") for d in data):
+            # If inputs is a list of objects with .pixel_values, concatenate them
+            pixel_values = [im.pixel_values if hasattr(im, "pixel_values") else None for im in data]
 
         page_table = kwargs.get("page_table", None)
         kv_cache = kwargs.get("kv_cache", None)
         vision_images = pixel_values
 
-        vision_images = [vision_images] if vision_images is not None else None
-
         return super().prefill_forward_text(
             tokens=inputs.input_ids,
             page_table=page_table,
@@ -482,3 +480,37 @@ def allocate_kv_cache(self, *args, **kwargs):
 
     def decode_forward(self, *args, **kwargs):
         return super().decode_forward_text(*args, **kwargs)
+
+
+class Gemma3ForCausalLM(Generator):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def initialize_vllm_model(
+        cls, hf_config, mesh_device, max_batch_size, max_seq_len=32768, n_layers=None, tt_data_parallel=1
+    ):
+        tt_model, model_args = initialize_vllm_text_transformer(
+            hf_config,
+            tt_data_parallel,
+            mesh_device,
+            max_batch_size,
+            max_seq_len=max_seq_len,
+            n_layers=n_layers,
+            dtype=ttnn.bfloat8_b,
+            optimizations=DecodersPrecision.performance,
+        )
+        return cls(tt_model, model_args, mesh_device)
+
+    @property
+    def cache_path(self):
+        return self.model_args[0].model_cache_path
+
+    def prefill_forward(self, *args, **kwargs):
+        return super().prefill_forward_text(*args, **kwargs)
+
+    def decode_forward(self, *args, **kwargs):
+        return super().decode_forward_text(*args, **kwargs)
+
+    def allocate_kv_cache(self, *args, **kwargs):
+        return allocate_vllm_kv_cache(*args, **kwargs, dp_model=self.model, tt_cache_path=self.cache_path)