@@ -1551,7 +1551,7 @@ class KVCacheConfigurator:
|
|||||||
# Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack.
|
# Mamba state is a fixed pre-capture allocation, so it can't ride the ~0 post-capture slack.
|
||||||
slack_gb = max(
|
slack_gb = max(
|
||||||
slack_gb,
|
slack_gb,
|
||||||
self.server_args.mamba_pre_capture_reserve_mb(
|
self.server_args.pre_capture_activation_reserve_mb(
|
||||||
get_device_memory_capacity(self.device)
|
get_device_memory_capacity(self.device)
|
||||||
)
|
)
|
||||||
/ 1024,
|
/ 1024,
|
||||||
|
|||||||
@@ -1722,6 +1722,8 @@ class MHATokenToKVPool(KVCache):
|
|||||||
quant_method=None,
|
quant_method=None,
|
||||||
post_capture_active: bool = False,
|
post_capture_active: bool = False,
|
||||||
):
|
):
|
||||||
|
self.k_buffer = None
|
||||||
|
self.v_buffer = None
|
||||||
if post_capture_active:
|
if post_capture_active:
|
||||||
# Reserved upper bound only (unbacked VA): page-align UP so
|
# Reserved upper bound only (unbacked VA): page-align UP so
|
||||||
# (size + page_size) % page_size == 0 holds for paged layouts.
|
# (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,
|
# 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
|
# e.g. vectorized_5d); fall back to _kv_buffer_shapes for the pre-allocation
|
||||||
# post-capture call, which only runs for NHD/HND.
|
# 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)
|
k_shape = tuple(self.k_buffer[0].shape)
|
||||||
v_shape = tuple(self.v_buffer[0].shape)
|
v_shape = tuple(self.v_buffer[0].shape)
|
||||||
else:
|
else:
|
||||||
@@ -3620,11 +3622,11 @@ class HybridLinearKVPool(KVCache):
|
|||||||
|
|
||||||
@property
|
@property
|
||||||
def post_capture_active(self) -> bool:
|
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
|
@property
|
||||||
def post_capture_backed_bytes(self) -> int:
|
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:
|
def finalize_backing(self, config) -> None:
|
||||||
# Only the attention KV is resized; the mamba state cache is fixed pre-capture.
|
# Only the attention KV is resized; the mamba state cache is fixed pre-capture.
|
||||||
|
|||||||
@@ -71,7 +71,7 @@ def compute_post_capture_kv_resize(
|
|||||||
if eager_decode_gap or mambaish_config(model_runner.model_config) is not None:
|
if eager_decode_gap or mambaish_config(model_runner.model_config) is not None:
|
||||||
headroom_gb = max(
|
headroom_gb = max(
|
||||||
headroom_gb,
|
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)
|
get_device_memory_capacity(model_runner.device)
|
||||||
)
|
)
|
||||||
/ 1024,
|
/ 1024,
|
||||||
|
|||||||
@@ -4747,9 +4747,9 @@ class ServerArgs:
|
|||||||
hf_config = self.get_model_config().hf_config
|
hf_config = self.get_model_config().hf_config
|
||||||
return not (is_deepseek_v4(hf_config) or is_minimax_sparse(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:
|
def pre_capture_activation_reserve_mb(self, gpu_mem: Optional[float]) -> float:
|
||||||
# Realistic runtime reserve for the fixed (non-resizable) mamba state cache,
|
# Runtime activation working-set reserve for eager decode above the captured
|
||||||
# which post-capture can't size from measured free memory.
|
# max_bs and transient prefill/logits; also covers fixed state caches.
|
||||||
if self.disaggregation_mode == "decode":
|
if self.disaggregation_mode == "decode":
|
||||||
running_requests = (
|
running_requests = (
|
||||||
self.max_running_requests or self.cuda_graph_config.decode.max_bs or 1
|
self.max_running_requests or self.cuda_graph_config.decode.max_bs or 1
|
||||||
|
|||||||
Reference in New Issue
Block a user