From 3c1717d9b6355e48855be7cc079db3ef9958f36d Mon Sep 17 00:00:00 2001
From: cctry <csycfl@gmail.com>
Date: Wed, 29 Jul 2026 15:03:59 -0700
Subject: [PATCH] Follow up on #30157 post-merge review (#32672)

---
 python/sglang/srt/mem_cache/kv_cache_configurator.py      | 2 +-
 python/sglang/srt/mem_cache/memory_pool.py                | 8 +++++---
 .../model_runner_components/kv_pool_runtime.py            | 2 +-
 python/sglang/srt/server_args.py                          | 6 +++---
 4 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py
index c30c474c7..bc4cf8154 100644
--- a/python/sglang/srt/mem_cache/kv_cache_configurator.py
+++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py
@@ -1551,7 +1551,7 @@ class KVCacheConfigurator:
             # Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack.
             slack_gb = max(
                 slack_gb,
-                self.server_args.mamba_pre_capture_reserve_mb(
+                self.server_args.pre_capture_activation_reserve_mb(
                     get_device_memory_capacity(self.device)
                 )
                 / 1024,
diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
index bfffed96c..000491f84 100644
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -1722,6 +1722,8 @@ class MHATokenToKVPool(KVCache):
         quant_method=None,
         post_capture_active: bool = False,
     ):
+        self.k_buffer = None
+        self.v_buffer = None
         if post_capture_active:
             # Reserved upper bound only (unbacked VA): page-align UP so
             # (size + page_size) % page_size == 0 holds for paged layouts.
@@ -2061,7 +2063,7 @@ class MHATokenToKVPool(KVCache):
         # Derive from the real buffers when they exist (covers arbitrary layouts,
         # e.g. vectorized_5d); fall back to _kv_buffer_shapes for the pre-allocation
         # post-capture call, which only runs for NHD/HND.
-        if getattr(self, "k_buffer", None) and getattr(self, "v_buffer", None):
+        if self.k_buffer and self.v_buffer:
             k_shape = tuple(self.k_buffer[0].shape)
             v_shape = tuple(self.v_buffer[0].shape)
         else:
@@ -3620,11 +3622,11 @@ class HybridLinearKVPool(KVCache):
 
     @property
     def post_capture_active(self) -> bool:
-        return getattr(self.full_kv_pool, "post_capture_active", False)
+        return self.full_kv_pool.post_capture_active
 
     @property
     def post_capture_backed_bytes(self) -> int:
-        return getattr(self.full_kv_pool, "post_capture_backed_bytes", 0)
+        return self.full_kv_pool.post_capture_backed_bytes
 
     def finalize_backing(self, config) -> None:
         # Only the attention KV is resized; the mamba state cache is fixed pre-capture.
diff --git a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py b/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py
index f72f76acd..96756e067 100644
--- a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py
+++ b/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py
@@ -71,7 +71,7 @@ def compute_post_capture_kv_resize(
     if eager_decode_gap or mambaish_config(model_runner.model_config) is not None:
         headroom_gb = max(
             headroom_gb,
-            model_runner.server_args.mamba_pre_capture_reserve_mb(
+            model_runner.server_args.pre_capture_activation_reserve_mb(
                 get_device_memory_capacity(model_runner.device)
             )
             / 1024,
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 35525e372..18b1d87e9 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -4747,9 +4747,9 @@ class ServerArgs:
         hf_config = self.get_model_config().hf_config
         return not (is_deepseek_v4(hf_config) or is_minimax_sparse(hf_config))
 
-    def mamba_pre_capture_reserve_mb(self, gpu_mem: Optional[float]) -> float:
-        # Realistic runtime reserve for the fixed (non-resizable) mamba state cache,
-        # which post-capture can't size from measured free memory.
+    def pre_capture_activation_reserve_mb(self, gpu_mem: Optional[float]) -> float:
+        # Runtime activation working-set reserve for eager decode above the captured
+        # max_bs and transient prefill/logits; also covers fixed state caches.
         if self.disaggregation_mode == "decode":
             running_requests = (
                 self.max_running_requests or self.cuda_graph_config.decode.max_bs or 1