From cd42dc6d4f876f9cbdfd7929e8eb63b5d7ccd9ca Mon Sep 17 00:00:00 2001 From: Matt Schwartz <113807100+matthematics1137@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:29:18 -0700 Subject: [PATCH] Log the recurrent-state slot reservation at model load For recurrent models with a draft model, every batch slot holds draft_num_tokens + 1 copies of the recurrent state, allocated on the GPU at load (728 MiB per slot for a 27B hybrid at draft length 4, 2.9 GiB at the default max_batch_size of 4). Log the size once the cache is built and say in the max_batch_size docs that single-user recurrent setups should set 1. The default is unchanged. Co-Authored-By: Claude Fable 5.1 --- backends/exllamav3/model.py | 28 ++++++++++++++++++++++++++++ common/config_models.py | 6 +++++- config_sample.yml | 4 ++++ docs/02.-Server-options.md | 2 +- 4 files changed, 38 insertions(+), 2 deletions(-) diff --git a/backends/exllamav3/model.py b/backends/exllamav3/model.py index 57e750bb..87f5ac80 100644 --- a/backends/exllamav3/model.py +++ b/backends/exllamav3/model.py @@ -405,6 +405,7 @@ async def create(cls, model_directory: pathlib.Path, hf_model: HFModel, **kwargs cache_mode_default = "FP16" self.cache_mode = unwrap(kwargs.get("cache_mode"), cache_mode_default) self.cache = self.create_cache(self.cache_mode, self.model) + self.log_recurrent_slot_cost() # Draft cache if self.use_draft_model: @@ -594,6 +595,33 @@ def job_max_rq_tokens(self, max_tokens: int) -> Optional[int]: return self.max_rq_tokens + def log_recurrent_slot_cost(self): + """ + Logs the VRAM reserved for recurrent state slots when drafting is enabled. + + Every batch slot of a recurrent (linear or sliding attention) model holds + max_history + 1 copies of the layer states, where max_history is the draft + length, so the allocation multiplies with both max_batch_size and + draft_num_tokens. It is reserved at load, before any request arrives. + """ + + recurrent_layers = getattr(self.cache, "recurrent_layers", None) + if not recurrent_layers or not self.cache.max_history: + return + + num_slots = self.cache.num_slots + num_states = self.cache.max_history + 1 + total_bytes = sum(layer.storage_size() for layer in recurrent_layers.values()) + state_bytes = sum(layer.get_checkpoint_size() for layer in recurrent_layers.values()) + slot_word = "slot" if num_slots == 1 else "slots" + hint = " Single-user setups can set max_batch_size: 1." if num_slots > 1 else "" + xlogger.info( + f"Recurrent state history: {num_slots} {slot_word} x {num_states} states " + f"({self.cache.max_history} draft tokens + 1) of {state_bytes / 1024**2:.0f} MiB " + f"each, {total_bytes / 1024**2:.0f} MiB of VRAM " + f"({total_bytes / num_slots / 1024**2:.0f} MiB per slot).{hint}" + ) + def create_cache(self, raw_cache_mode: str, model: Model): # Cast exl2 types to exl3 match raw_cache_mode: diff --git a/common/config_models.py b/common/config_models.py index 596764c7..1fc0fa39 100644 --- a/common/config_models.py +++ b/common/config_models.py @@ -369,7 +369,11 @@ class ModelConfig(BaseConfigModel): "The default maximum batch size for transformer architectures is 32. Recurrent\n" "models with linear or sliding attention use more VRAM to support larger batches,\n" "so the default value is reduced to 4. If you do not require concurrency at all, you\n" - "can reduce it further to minimize VRAM overhead." + "can reduce it further to minimize VRAM overhead.\n" + "With a draft model, every slot of a recurrent model reserves draft_num_tokens + 1\n" + "copies of the recurrent state at load (728 MiB per slot for a 27B hybrid at the\n" + "default draft length of 4; the total is logged at load), so single-user setups\n" + "serving a recurrent model should set this to 1." ), ge=1, ) diff --git a/config_sample.yml b/config_sample.yml index 620a8943..1d1a61fd 100644 --- a/config_sample.yml +++ b/config_sample.yml @@ -201,6 +201,10 @@ model: # models with linear or sliding attention use more VRAM to support larger batches, # so the default value is reduced to 4. If you do not require concurrency at all, you # can reduce it further to minimize VRAM overhead. + # With a draft model, every slot of a recurrent model reserves draft_num_tokens + 1 + # copies of the recurrent state at load (728 MiB per slot for a 27B hybrid at the + # default draft length of 4; the total is logged at load), so single-user setups + # serving a recurrent model should set this to 1. max_batch_size: # Set the prompt template for this model. (default: None) diff --git a/docs/02.-Server-options.md b/docs/02.-Server-options.md index 8f520dd6..3599b456 100644 --- a/docs/02.-Server-options.md +++ b/docs/02.-Server-options.md @@ -84,7 +84,7 @@ Note: Most of the options here will only apply on initial model load/startup (ep | cache_mode | String ("FP16") | Cache mode for the model. Specify the pair `k_bits,v_bits` where each is an integer from 2-8 (e.g. `8,8`).

The legacy values FP16, Q8, Q6 and Q4 are also accepted. | | cache_size | Int (max_seq_len) | Size of the K/V cache

Note: If using CFG, the cache size should be 2 * max_seq_len. | | chunk_size | Int (2048) | Amount of tokens per chunk with ingestion. A lower value reduces VRAM usage at the cost of ingestion speed. | -| max_batch_size | Int (None) | The absolute maximum amount of prompts to process at one time. This value is automatically adjusted based on cache size. | +| max_batch_size | Int (None) | The absolute maximum amount of prompts to process at one time. This value is automatically adjusted based on cache size.

Note: With a draft model, every slot of a recurrent (linear or sliding attention) model reserves `draft_num_tokens + 1` copies of the recurrent state at load (728 MiB per slot for a 27B hybrid at the default draft length of 4; the total is logged at load). Single-user setups serving a recurrent model should set this to 1. | | prompt_template | String (None) | Name of a jinja2 chat template to apply for this model. Must be located in the `templates` directory. | | vision | Bool (False) | Enable vision support for the provided model (if it exists). | | vision_offload | Bool (False) | Keep the vision model's weights in pinned system RAM and stream them to the GPU during inference, trading vision speed for VRAM. Only applies when vision is enabled. |