Make mem_fraction_static reserve disaggregation-mode aware (#29615)
This commit is contained in:
@@ -3529,17 +3529,28 @@ class ServerArgs:
|
|||||||
if self.mem_fraction_static is None:
|
if self.mem_fraction_static is None:
|
||||||
# Constant meta data (e.g., from attention backend)
|
# Constant meta data (e.g., from attention backend)
|
||||||
reserved_mem = 512
|
reserved_mem = 512
|
||||||
# For activation during large prefill
|
# For activation slack
|
||||||
if self.chunked_prefill_size > 0:
|
if self.disaggregation_mode == "decode":
|
||||||
|
# Decode nodes do no prefill; size activation to the decode batch.
|
||||||
|
running_requests = (
|
||||||
|
self.max_running_requests or decode_cuda_graph_config.max_bs or 1
|
||||||
|
)
|
||||||
|
draft_tokens = self.speculative_num_draft_tokens or 1
|
||||||
|
reserved_mem += max(running_requests * draft_tokens, 2048) * 1.5
|
||||||
|
elif self.chunked_prefill_size > 0:
|
||||||
reserved_mem += max(self.chunked_prefill_size, 2048) * 1.5
|
reserved_mem += max(self.chunked_prefill_size, 2048) * 1.5
|
||||||
else:
|
else:
|
||||||
reserved_mem += max(self.max_prefill_tokens, 2048) * 1.5
|
reserved_mem += max(self.max_prefill_tokens, 2048) * 1.5
|
||||||
# For cuda graphs
|
# For decode cuda graphs (skip on prefill-only nodes)
|
||||||
reserved_mem += decode_cuda_graph_config.max_bs * 2
|
if (
|
||||||
|
self.disaggregation_mode != "prefill"
|
||||||
|
and decode_cuda_graph_config.backend != Backend.DISABLED
|
||||||
|
):
|
||||||
|
reserved_mem += decode_cuda_graph_config.max_bs * 2
|
||||||
# Some adjustments for large parallel size
|
# Some adjustments for large parallel size
|
||||||
reserved_mem += self.tp_size * self.pp_size / 8 * 1024
|
reserved_mem += self.tp_size * self.pp_size / 8 * 1024
|
||||||
|
|
||||||
if self.enable_dp_attention:
|
if self.enable_dp_attention and self.disaggregation_mode != "prefill":
|
||||||
# DP attention needs more padding for some operations
|
# DP attention needs more padding for some operations
|
||||||
reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 3
|
reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 3
|
||||||
|
|
||||||
@@ -3549,8 +3560,11 @@ class ServerArgs:
|
|||||||
if decode_cuda_graph_config.max_bs > 300:
|
if decode_cuda_graph_config.max_bs > 300:
|
||||||
reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 1.5
|
reserved_mem += decode_cuda_graph_config.max_bs * self.dp_size * 1.5
|
||||||
|
|
||||||
# For piecewise cuda graphs
|
# For prefill piecewise cuda graphs (skip on decode-only nodes)
|
||||||
if prefill_cuda_graph_config.backend != Backend.DISABLED:
|
if (
|
||||||
|
self.disaggregation_mode != "decode"
|
||||||
|
and prefill_cuda_graph_config.backend != Backend.DISABLED
|
||||||
|
):
|
||||||
if not self.use_mla_backend():
|
if not self.use_mla_backend():
|
||||||
# Only calculate the memory overhead for Non-Torch Memory use since the Torch Memory can be reused with Cuda Graph Capture
|
# Only calculate the memory overhead for Non-Torch Memory use since the Torch Memory can be reused with Cuda Graph Capture
|
||||||
reserved_mem += len(prefill_cuda_graph_config.bs) * 8
|
reserved_mem += len(prefill_cuda_graph_config.bs) * 8
|
||||||
@@ -3561,6 +3575,15 @@ class ServerArgs:
|
|||||||
if gpu_mem is not None and gpu_mem > 60 * 1024:
|
if gpu_mem is not None and gpu_mem > 60 * 1024:
|
||||||
reserved_mem = max(reserved_mem, 10 * 1024)
|
reserved_mem = max(reserved_mem, 10 * 1024)
|
||||||
|
|
||||||
|
# DeepEP all-to-all buffers captured in the decode graph are real
|
||||||
|
# extra allocations, so reserve them on top of the floor.
|
||||||
|
if (
|
||||||
|
self.disaggregation_mode != "prefill"
|
||||||
|
and decode_cuda_graph_config.backend != Backend.DISABLED
|
||||||
|
and self.moe_a2a_backend == "deepep"
|
||||||
|
):
|
||||||
|
reserved_mem += 2 * 1024
|
||||||
|
|
||||||
self.mem_fraction_static = (
|
self.mem_fraction_static = (
|
||||||
round((gpu_mem - reserved_mem) / gpu_mem, 3)
|
round((gpu_mem - reserved_mem) / gpu_mem, 3)
|
||||||
if gpu_mem is not None
|
if gpu_mem is not None
|
||||||
|
|||||||
Reference in New Issue
Block a user