From 9509c4ccd3977d85f89d0fd15950c4dc72b28108 Mon Sep 17 00:00:00 2001
From: Zilin Zhu <zhuzilinallen@gmail.com>
Date: Sun, 16 Nov 2025 17:12:52 +0800
Subject: [PATCH] [RL] enable offloading hybrid linear attn model (#13336)

---
 python/sglang/srt/mem_cache/memory_pool.py    | 19 ++++++++++++-------
 .../sglang/srt/model_executor/model_runner.py |  1 +
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
index 8cd3865b2..10f16dfa2 100644
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -366,12 +366,16 @@ class HybridReqToTokenPool(ReqToTokenPool):
             device=device,
             enable_memory_saver=enable_memory_saver,
         )
-        self._init_mamba_pool(
-            size=mamba_size,
-            cache_params=cache_params,
-            device=device,
-            speculative_num_draft_tokens=speculative_num_draft_tokens,
+        memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
         )
+        with memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            self._init_mamba_pool(
+                size=mamba_size,
+                cache_params=cache_params,
+                device=device,
+                speculative_num_draft_tokens=speculative_num_draft_tokens,
+            )
 
     def _init_mamba_pool(
         self,
@@ -948,6 +952,7 @@ class HybridLinearKVPool(KVCache):
         enable_kvcache_transpose: bool,
         device: str,
         mamba_pool: MambaPool,
+        enable_memory_saver: bool,
         # TODO: refactor mla related args
         use_mla: bool = False,
         kv_lora_rank: int = None,
@@ -979,7 +984,7 @@ class HybridLinearKVPool(KVCache):
                 head_dim=head_dim,
                 layer_num=self.full_layer_nums,
                 device=device,
-                enable_memory_saver=False,
+                enable_memory_saver=enable_memory_saver,
             )
         else:
             TokenToKVPoolClass = MLATokenToKVPool
@@ -991,7 +996,7 @@ class HybridLinearKVPool(KVCache):
                 device=device,
                 kv_lora_rank=kv_lora_rank,
                 qk_rope_head_dim=qk_rope_head_dim,
-                enable_memory_saver=False,
+                enable_memory_saver=enable_memory_saver,
             )
         self.full_attention_layer_id_mapping = {
             id: i for i, id in enumerate(full_attention_layer_ids)
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index a7f0f498f..fb3dbd4be 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -1837,6 +1837,7 @@ class ModelRunner:
                     enable_kvcache_transpose=False,
                     device=self.device,
                     mamba_pool=self.req_to_token_pool.mamba_pool,
+                    enable_memory_saver=self.server_args.enable_memory_saver,
                     use_mla=self.use_mla_backend,
                     **extra_args,
                 )