Clean up page size padding logic

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>

Author: tdoublep

vllm/v1/kv_cache_interface.py

@@ -11,7 +11,7 @@
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.utils import cdiv, get_dtype_size, round_up
+from vllm.utils import cdiv, get_dtype_size
 
 logger = init_logger(__name__)
 
@@ -159,7 +159,7 @@ def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
 class MambaSpec(KVCacheSpec):
     shapes: tuple[tuple[int, ...], ...]
     dtype: torch.dtype
-    multiple_of: Optional[int]
+    page_size_padded: Optional[int] = None
 
     def __post_init__(self):
         self.num_elements = sum(prod(shape) for shape in self.shapes)
@@ -171,8 +171,9 @@ def type_id(self) -> str:
     @property
     def page_size_bytes(self) -> int:
         page_size = self.num_elements * get_dtype_size(self.dtype)
-        if self.multiple_of is not None:
-            page_size = round_up(page_size, self.multiple_of)
+        if self.page_size_padded is not None:
+            assert self.page_size_padded >= page_size
+            return self.page_size_padded
         return page_size
 
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:

vllm/v1/worker/gpu_model_runner.py

@@ -2575,15 +2575,6 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
         mamba_layers = get_layers_from_vllm_config(self.vllm_config,
                                                    MambaMixer2)
         if len(mamba_layers) > 0:
-            if len(attn_layers) > 0:
-                # Mamba state must be padded to an integer number of
-                # 16th tokens worth of attention pages
-                attn_layer_name = next(iter(attn_layers))
-                attn_page_size = kv_cache_spec[attn_layer_name].page_size_bytes
-                multiple_of = 16 * attn_page_size // block_size
-            else:
-                multiple_of = None
-
             if self.vllm_config.speculative_config is not None:
                 raise NotImplementedError(
                     "Mamba with speculative decoding is not supported yet.")
@@ -2594,25 +2585,39 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
                 raise NotImplementedError(
                     "Prefix caching is not supported for Mamba yet.")
             max_model_len = self.vllm_config.model_config.max_model_len
+
+            if len(attn_layers) > 0:
+                attn_layer_name = next(iter(attn_layers))
+                attn_page_size = kv_cache_spec[attn_layer_name].page_size_bytes
+                mamba_layer_name = next(iter(mamba_layers))
+                mamba_page_size = MambaSpec(
+                    shapes=mamba_layers[mamba_layer_name].get_state_shape(),
+                    dtype=self.kv_cache_dtype,
+                    block_size=max_model_len).page_size_bytes
+                if attn_page_size < mamba_page_size:
+                    # attention page size (for 16 tokens)
+                    attn_page_size_16 = 16 * attn_page_size // block_size
+                    # some attention backends (e.g. FA) only support setting
+                    # block size to multiple of 16, so let's suggest a value
+                    # that would work (note: FA is currently not compatible
+                    # with mamba layers, use FlashInfer instead).
+                    suggest_attn_block_size = 16 * cdiv(
+                        mamba_page_size, attn_page_size_16)
+                    raise ValueError(
+                        "Attention block size should be increased to at least "
+                        f"{suggest_attn_block_size} in order to match "
+                        "the mamba page size")
+                page_size_padded = attn_page_size
+            else:
+                page_size_padded = None
+
             # Set block_size to max_model_len, so that mamba model will always
             # have only one block in the KV cache.
             for layer_name, mamba_module in mamba_layers.items():
                 kv_cache_spec[layer_name] = MambaSpec(
                     shapes=mamba_module.get_state_shape(),
                     dtype=self.kv_cache_dtype,
                     block_size=max_model_len,
-                    multiple_of=multiple_of)
-
-            if len(attn_layers) > 0:
-                mamba_layer_name = next(iter(mamba_layers))
-                mamba_page_size = kv_cache_spec[
-                    mamba_layer_name].page_size_bytes
-                if attn_page_size < mamba_page_size:
-                    required_attn_block_size = cdiv(mamba_page_size,
-                                                    multiple_of) * 16
-                    raise ValueError(
-                        "Attention block size must be increased to "
-                        f"{required_attn_block_size} in order to match "
-                        "the mamba page size")
+                    page_size_padded=page_size_padded)
 
         return kv_cache_spec