fix: set num_blocks outside of model_runner (#410)

Author: rebel-eunji

vllm_rbln/model_executor/models/optimum/model_base.py

@@ -67,8 +67,13 @@ def _env_int(name: str, default: int) -> int:
 
     def _estimated_num_blocks(self) -> int:
         """Override estimated blocks if num_gpu_blocks_override is set."""
-        num_gpu_blocks_override = self.vllm_config.cache_config.num_gpu_blocks_override
-        if num_gpu_blocks_override is not None:
+        if (
+            self.vllm_config.additional_config
+            and "num_blocks_override" in self.vllm_config.additional_config
+        ):
+            num_gpu_blocks_override = self.vllm_config.additional_config[
+                "num_blocks_override"
+            ]
             return num_gpu_blocks_override
         else:
             return int(self.estimated_kvcache_num_blocks)

vllm_rbln/utils/optimum/configuration.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import math
 import os
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -32,9 +33,34 @@
 logger = init_logger(__name__)
 
 
+def is_full_block_available(num_blocks: int, vllm_config: VllmConfig) -> bool:
+    if vllm_config.cache_config.enable_prefix_caching:
+        block_size = vllm_config.additional_config["attn_block_size"]
+
+    else:
+        block_size = vllm_config.cache_config.block_size
+
+    max_model_len = vllm_config.model_config.max_model_len
+    max_num_seqs = vllm_config.scheduler_config.max_num_seqs
+
+    blocks_per_seq = math.ceil(max_model_len / block_size)
+    ideal_total = max_num_seqs * blocks_per_seq
+    return num_blocks >= ideal_total
+
+
+def get_block_ratio(vllm_config: VllmConfig) -> int:
+    if vllm_config.cache_config.enable_prefix_caching:
+        ob_size = vllm_config.additional_config["attn_block_size"]
+        ib_size = vllm_config.cache_config.block_size
+        blk_ratio = ob_size // ib_size
+    else:
+        blk_ratio = 1
+    return blk_ratio
+
+
 def get_rbln_params(
     vllm_config: VllmConfig, rbln_config: dict
-) -> tuple[int, int, int, int]:
+) -> tuple[int, int, int, int, int]:
     kvcache_block_size = None
     prefill_chunk_size = 128
     batch_size = None
@@ -44,11 +70,13 @@ def get_rbln_params(
         max_seq_len = rbln_config.get("dec_max_seq_len")
         kvcache_block_size = max_seq_len
         batch_size = rbln_config.get("batch_size")
+        num_blocks = rbln_config.get("kvcache_num_blocks")
     elif is_multi_modal(vllm_config.model_config.hf_config):
         # Get configurations from main module (e.g. Qwen2.5-VL, Whisper)
         kvcache_block_size = rbln_config.get("kvcache_block_size")
         batch_size = rbln_config.get("batch_size")
         max_seq_len = rbln_config.get("max_seq_len")
+        num_blocks = rbln_config.get("kvcache_num_blocks")
         if max_seq_len is None:  # Whisper FIXME to be moved to enc-dec
             max_seq_len = rbln_config.get("dec_max_seq_len")
         # Get configurations from submodule
@@ -61,19 +89,26 @@ def get_rbln_params(
                     )
                     batch_size = rbln_config[submodule].get("batch_size", None)
                     max_seq_len = rbln_config[submodule].get("max_seq_len", None)
+                    num_blocks = rbln_config[submodule].get("kvcache_num_blocks", None)
                     if kvcache_block_size is not None:
                         break
 
     elif is_pooling_arch(vllm_config.model_config.hf_config):
         max_seq_len = rbln_config.get("max_seq_len")
         kvcache_block_size = max_seq_len
         batch_size = rbln_config.get("batch_size")
+        num_blocks = rbln_config.get("kvcache_num_blocks")
+        if num_blocks is None:
+            num_blocks = batch_size  # for pooling models, each sequence is one block
     else:
         # decoder
         kvcache_block_size = rbln_config.get("kvcache_block_size")
         prefill_chunk_size = rbln_config.get("prefill_chunk_size", 128)
         batch_size = rbln_config.get("batch_size")
         max_seq_len = rbln_config.get("max_seq_len")
+        num_blocks = rbln_config.get("kvcache_num_blocks")
+
+    assert num_blocks is not None, "num_blocks must be specified in rbln_config.json"
 
     assert kvcache_block_size is not None, (
         "kvcache_block_size must be specified in rbln_config.json"
@@ -83,7 +118,7 @@ def get_rbln_params(
     # NOTE:
     # prefill_chunk_size is only used for decoder-only models
     # with prefix caching
-    return kvcache_block_size, batch_size, max_seq_len, prefill_chunk_size
+    return num_blocks, batch_size, max_seq_len, kvcache_block_size, prefill_chunk_size
 
 
 def set_block_size_for_prefix_caching(
@@ -132,6 +167,7 @@ def set_block_size_for_prefix_caching(
 
 def update_vllm_config_with_rbln_params(
     vllm_config: VllmConfig,
+    num_blocks: int,
     batch_size: int,
     max_model_len: int,
     kvcache_block_size: int,
@@ -181,6 +217,23 @@ def update_vllm_config_with_rbln_params(
             )
             vllm_config.cache_config.block_size = kvcache_block_size
 
+    # num_blocks is determined by rbln_config or overridden by user.
+    if vllm_config.cache_config.num_gpu_blocks_override is not None:
+        num_blocks = vllm_config.cache_config.num_gpu_blocks_override
+        vllm_config.additional_config["num_blocks_override"] = num_blocks
+
+    blk_ratio = get_block_ratio(vllm_config)
+
+    if is_full_block_available(num_blocks, vllm_config):
+        adjusted_num_blocks = num_blocks * blk_ratio + 1
+    else:
+        adjusted_num_blocks = (num_blocks - 1) * blk_ratio + 1
+
+    vllm_config.cache_config.num_blocks = adjusted_num_blocks
+
+    if vllm_config.cache_config.num_gpu_blocks_override is not None:
+        vllm_config.cache_config.num_gpu_blocks_override = adjusted_num_blocks
+
 
 def is_qwen3_pooling(
     vllm_config: VllmConfig,
@@ -215,11 +268,16 @@ def sync_with_rbln_config(vllm_config: VllmConfig) -> None:
         raise RuntimeError("Failed to get RBLN config: %s", e) from e
 
     if rbln_config is not None:
-        kvcache_block_size, batch_size, max_model_len, prefill_chunk_size = (
-            get_rbln_params(vllm_config, rbln_config)
-        )
+        (
+            num_blocks,
+            batch_size,
+            max_model_len,
+            kvcache_block_size,
+            prefill_chunk_size,
+        ) = get_rbln_params(vllm_config, rbln_config)
         update_vllm_config_with_rbln_params(
             vllm_config,
+            num_blocks,
             batch_size,
             max_model_len,
             kvcache_block_size,

vllm_rbln/v1/worker/optimum_worker.py

@@ -112,6 +112,12 @@ def determine_available_memory(self) -> int:
 
         adapter = self.model_runner.model.kv_block_adapter
         num_gpu_blocks = adapter.get_available_num_blocks()
+        validation_blocks = self.model_runner.vllm_config.cache_config.num_blocks
+        # This will be removed after validation check
+        assert num_gpu_blocks == validation_blocks, (
+            f"The number of blocks from the model runner ({num_gpu_blocks}) "
+            f"and the platform ({validation_blocks}) must be the same."
+        )
 
         return num_gpu_blocks * page_size * num_layers