From 8abf805de2c948604f7dcf82b61abae89eb96fb0 Mon Sep 17 00:00:00 2001
From: Tsury <4947581+Tsury@users.noreply.github.com>
Date: Sat, 17 Oct 2026 21:55:49 +0300
Subject: [PATCH] feat: expose fp32 recurrent state storage

---
 backends/exllamav3/model.py   | 4 ++++
 common/config_models.py       | 7 +++++++
 endpoints/core/types/model.py | 1 +
 3 files changed, 12 insertions(+)

diff --git a/backends/exllamav3/model.py b/backends/exllamav3/model.py
index c1d5074d..f420c036 100644
--- a/backends/exllamav3/model.py
+++ b/backends/exllamav3/model.py
@@ -253,6 +253,7 @@ class ExllamaV3Container:
     max_seq_len: int = 4096
     cache_size: int = 4096
     cache_mode: str = "FP16"
+    recurrent_state_fp32: bool = False
     draft_cache_mode: str = "FP16"
     chunk_size: int = 2048
     max_rq_tokens: Optional[int] = 2048
@@ -607,6 +608,8 @@ async def create(cls, model_directory: pathlib.Path, hf_model: HFModel, **kwargs
         # Create cache
         cache_mode_default = "FP16"
         self.cache_mode = unwrap(kwargs.get("cache_mode"), cache_mode_default)
+        self.recurrent_state_fp32 = unwrap(kwargs.get("recurrent_state_fp32"), False)
+        xlogger.info(f"Recurrent state storage: {'fp32' if self.recurrent_state_fp32 else 'fp16'}")
         self.cache = self.create_cache(self.cache_mode, self.model)
 
         # Draft cache
@@ -901,6 +904,7 @@ def create_cache(self, raw_cache_mode: str, model: Model):
         batch_draft_args = {
             "max_batch_size": self.max_batch_size,
             "max_history": max(draft_tokens, self.hybrid_draft_tokens),
+            "recurrent_dtype": torch.float if self.recurrent_state_fp32 else torch.half,
         }
 
         if split_cache_mode:
diff --git a/common/config_models.py b/common/config_models.py
index 4aa37dec..77c9653c 100644
--- a/common/config_models.py
+++ b/common/config_models.py
@@ -257,6 +257,13 @@ class ModelConfig(BaseConfigModel):
             f"The legacy values {str(CACHE_SIZES)[15:-1]} are also accepted."
         ),
     )
+    recurrent_state_fp32: Optional[bool] = Field(
+        False,
+        description=(
+            "Store GatedDeltaNet/KDA recurrent state in FP32 instead of FP16. "
+            "Uses twice the state and checkpoint memory; scans compute in FP32 either way."
+        ),
+    )
     tensor_parallel: Optional[bool] = Field(
         False,
         description=(
diff --git a/endpoints/core/types/model.py b/endpoints/core/types/model.py
index 62cf23e4..057d3ba0 100644
--- a/endpoints/core/types/model.py
+++ b/endpoints/core/types/model.py
@@ -147,6 +147,7 @@ class ModelLoadRequest(_WarnOnUnknownFields):
         examples=[4096],
     )
     cache_mode: Optional[str] = None
+    recurrent_state_fp32: Optional[bool] = None
     tensor_parallel: Optional[bool] = None
     tensor_parallel_backend: Optional[str] = "native"
     gpu_split_auto: Optional[bool] = None