Expose native FP32 recurrent state storage for ExLlamaV3

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:

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=(

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