Fixes: gemma3's block_table scheduling bug due to padding (#28)

* fix shape of pixel_values

* pre-commit

* forward with padded

* change the parameter and register gemma3

* allow kwargs for register vllm model

* pre-commit

* import error

* revert examples

* RBLNGemma3MultiModalProcessor's location

* pre-commit

* clean

* clean

* pre-commit

* pre-commit

* pre-commit

* fix param

* fix pre-commit

* fixme comment

* fix pre-commit

* rebase

* fix typo

---------

Co-authored-by: rebel-eunji <eunji.lee@rebellions.ai>

Author: rebel-thkim

vllm_rbln/__init__.py

@@ -30,3 +30,7 @@ def register_model():
     )
     ModelRegistry.register_model("T5EncoderModel",
                                  "optimum.rbln:RBLNT5EncoderModel")
+    ModelRegistry.register_model(
+        "Gemma3ForConditionalGeneration",
+        "vllm_rbln.model_executor.models.optimum.gemma3:RBLNOptimumGemma3ForConditionalGeneration"
+    )

vllm_rbln/model_executor/model_loader/rbln_model_loader.py

@@ -12,13 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch.nn as nn
-from vllm.config import ModelConfig, SchedulerConfig
+from vllm.config import VllmConfig
 
 from vllm_rbln.model_executor.models.optimum import load_model
 
 
-def get_optimum_model(
-    model_config: ModelConfig,
-    scheduler_config: SchedulerConfig,
-) -> nn.Module:
-    return load_model(model_config, scheduler_config)
+def get_optimum_model(vllm_config: VllmConfig, ) -> nn.Module:
+    return load_model(vllm_config)

vllm_rbln/model_executor/models/optimum/__init__.py

@@ -14,7 +14,7 @@
 """Utilities for selecting and loading rbln models."""
 
 import torch.nn as nn
-from vllm.config import ModelConfig, SchedulerConfig
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
 
 from .base import (_RBLN_MULTIMODAL_MODELS, ModelInputForRBLN,
@@ -41,31 +41,26 @@
 }
 
 
-def load_model(
-    model_config: ModelConfig,
-    scheduler_config: SchedulerConfig,
-) -> nn.Module:
+def load_model(vllm_config: VllmConfig) -> nn.Module:
+    model_config = vllm_config.model_config
+
     if is_multi_modal(model_config.hf_config):
         architectures = getattr(model_config.hf_config, "architectures", [])
         if architectures[0] in _RBLN_OPTIMUM_MULTIMODAL_MODELS:
             rbln_model_arch = _RBLN_OPTIMUM_MULTIMODAL_MODELS[architectures[0]]
-            rbln_model = rbln_model_arch(model_config=model_config,
-                                         scheduler_config=scheduler_config)
+            rbln_model = rbln_model_arch(vllm_config)
         else:
             raise NotImplementedError(
                 f"Model architectures {architectures} are "
                 f"not supported on RBLN Optimum for now. "
                 "Supported multimodal architectures: "
                 f"{list(_RBLN_OPTIMUM_MULTIMODAL_MODELS.keys())}")
     elif is_enc_dec_arch(model_config.hf_config):
-        rbln_model = RBLNOptimumEncoderDecoder(
-            model_config=model_config, scheduler_config=scheduler_config)
+        rbln_model = RBLNOptimumEncoderDecoder(vllm_config)
     elif is_pooling_arch(model_config.hf_config):
-        rbln_model = RBLNOptimumForEncoderModel(
-            model_config=model_config, scheduler_config=scheduler_config)
+        rbln_model = RBLNOptimumForEncoderModel(vllm_config)
     else:
-        rbln_model = RBLNOptimumForCausalLM(model_config=model_config,
-                                            scheduler_config=scheduler_config)
+        rbln_model = RBLNOptimumForCausalLM(vllm_config)
     return rbln_model.eval()
 
 

vllm_rbln/model_executor/models/optimum/blip2.py

@@ -14,7 +14,7 @@
 from typing import Any, Optional
 
 import torch
-from vllm.config import ModelConfig, SchedulerConfig
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.models.blip2 import (Blip2ImageEmbeddingInputs,
                                               Blip2ImageInputs,
@@ -32,23 +32,22 @@ class RBLNOptimumBlip2ForConditionalGeneration(RBLNOptimumModelBase,
 
     def __init__(
         self,
-        model_config: ModelConfig,
-        scheduler_config: SchedulerConfig,
+        vllm_config: VllmConfig,
     ) -> None:
-        super().__init__(model_config=model_config,
-                         scheduler_config=scheduler_config)
+        super().__init__(vllm_config=vllm_config)
         self.setup_decoder_mixin(
             attn_impl=self.attn_impl,
             padding_value=self.padding_value,
-            vocab_size=model_config.get_vocab_size,
+            vocab_size=self.model_config.get_vocab_size,
             use_multiple_decoder=getattr(self.model.rbln_config.language_model,
                                          "use_multiple_decoder", False),
             default_batch_size=self.scheduler_config.max_num_seqs,
             decoder_batch_sizes=self.model.rbln_config.language_model.
             decoder_batch_sizes,
         )
 
-    def forward(self, model_input: ModelInputForRBLN) -> torch.Tensor:
+    def forward(self, model_input: ModelInputForRBLN,
+                **kwargs) -> torch.Tensor:
         input_ids = model_input.input_tokens
         cache_position = model_input.input_positions
         block_tables = model_input.block_tables

vllm_rbln/model_executor/models/optimum/decoder_only.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
-from vllm.config import ModelConfig, SchedulerConfig
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
 
 from .base import ModelInputForRBLN, version_error
@@ -25,25 +25,21 @@ class RBLNOptimumForCausalLM(RBLNOptimumModelBase, RBLNOptimumDecoderMixin):
 
     def __init__(
         self,
-        model_config: ModelConfig,
-        scheduler_config: SchedulerConfig,
-        **kwargs,
+        vllm_config: VllmConfig,
     ) -> None:
-        super().__init__(
-            model_config=model_config,
-            scheduler_config=scheduler_config,
-        )
+        super().__init__(vllm_config=vllm_config)
         self.setup_decoder_mixin(
             attn_impl=self.attn_impl,
             padding_value=self.padding_value,
-            vocab_size=model_config.get_vocab_size,
+            vocab_size=self.model_config.get_vocab_size,
             use_multiple_decoder=getattr(self.model.rbln_config,
                                          "use_multiple_decoder", False),
             default_batch_size=self.scheduler_config.max_num_seqs,
             decoder_batch_sizes=self.model.rbln_config.decoder_batch_sizes,
         )
 
-    def forward(self, model_input: ModelInputForRBLN) -> torch.Tensor:
+    def forward(self, model_input: ModelInputForRBLN,
+                **kwargs) -> torch.Tensor:
         input_ids = model_input.input_tokens
         cache_position = model_input.input_positions
         block_tables = model_input.block_tables

vllm_rbln/model_executor/models/optimum/encoder.py

@@ -15,7 +15,7 @@
 from typing import Optional
 
 import torch
-from vllm.config import ModelConfig, PoolerConfig, SchedulerConfig
+from vllm.config import PoolerConfig, VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.pooler import Pooler, PoolingType
 from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
@@ -31,11 +31,10 @@ class RBLNOptimumForEncoderModel(RBLNOptimumModelBase):
 
     def __init__(
         self,
-        model_config: ModelConfig,
-        scheduler_config: SchedulerConfig,
+        vllm_config: VllmConfig,
     ) -> None:
-        super().__init__(model_config, scheduler_config)
-        self._pooler = self._build_pooler(model_config.pooler_config)
+        super().__init__(vllm_config=vllm_config)
+        self._pooler = self._build_pooler(self.model_config.pooler_config)
 
     def is_classification_arch(self):
         architectures = getattr(
@@ -94,7 +93,8 @@ def _build_pooler(self, pooler_config: PoolerConfig) -> Optional[Pooler]:
             )
         return None
 
-    def forward(self, model_input: ModelInputForRBLN) -> torch.Tensor:
+    def forward(self, model_input: ModelInputForRBLN,
+                **kwargs) -> torch.Tensor:
         input_ids, token_type_ids, positions = self.preprocess(
             model_input.input_tokens,
             model_input.token_type_ids,

vllm_rbln/model_executor/models/optimum/encoder_decoder.py

@@ -14,7 +14,7 @@
 from typing import List, Optional, Union
 
 import torch
-from vllm.config import ModelConfig, SchedulerConfig
+from vllm.config import VllmConfig
 from vllm.logger import init_logger
 
 from .base import ModelInputForRBLN, version_error
@@ -28,17 +28,15 @@ class RBLNOptimumEncoderDecoder(RBLNOptimumModelBase, RBLNOptimumDecoderMixin):
 
     def __init__(
         self,
-        model_config: ModelConfig,
-        scheduler_config: SchedulerConfig,
+        vllm_config: VllmConfig,
     ) -> None:
-        super().__init__(model_config=model_config,
-                         scheduler_config=scheduler_config)
+        super().__init__(vllm_config=vllm_config)
         # encoder length used for encoder_decoder architecture
         self.enc_lengths = [0] * self.batch_size
         self.setup_decoder_mixin(
             attn_impl=self.attn_impl,
             padding_value=self.padding_value,
-            vocab_size=model_config.get_vocab_size,
+            vocab_size=self.model_config.get_vocab_size,
             use_multiple_decoder=False,
             default_batch_size=self.scheduler_config.max_num_seqs,
             decoder_batch_sizes=[self.batch_size],
@@ -115,7 +113,8 @@ def _forward(
 
         return logits
 
-    def forward(self, model_input: ModelInputForRBLN) -> torch.Tensor:
+    def forward(self, model_input: ModelInputForRBLN,
+                **kwargs) -> torch.Tensor:
         input_ids = model_input.input_tokens
         cache_position = model_input.input_positions
         is_prompt = model_input.sampling_metadata.num_prompts > 0