fix: raise MemoryError when available_dram becomes negative (#80)

Author: junstar92

vllm_rbln/worker/utils.py

@@ -78,6 +78,10 @@ def align_2MB(x: int) -> int:
     available_dram = tensor_parallel_size * (ATOM_DRAM_NBYTES -
                                              ATOM_SYS_DRAM_NBYTES)
 
+    def check_oom(available_dram: int) -> None:
+        if available_dram <= 0:
+            raise MemoryError("Insufficient DRAM during block calculation.")
+
     if kernel_size is None:
         if n_model_params is None:
             raise ValueError("`n_model_params` should be specified \
@@ -106,6 +110,8 @@ def align_2MB(x: int) -> int:
         buffer = buffer_per_core * tensor_parallel_size
     available_dram -= buffer
 
+    check_oom(available_dram)
+
     b = kvcache_block_size * align(head_dim, 64) * math.ceil(
         num_key_value_heads / tensor_parallel_size) * 2
     c = num_layers * 2 * tensor_parallel_size