karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[copy of #45208] CuMem slept-L1 fragmentation accounting (#49208)

Browse Source 
Signed-off-by: haosdent <haosdent@gmail.com>
Signed-off-by: Justin Wood <justin.m.wood@me.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: haosdent <haosdent@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Justin Wood <jwood@me.com>
This commit is contained in:
Matthew Bonanni
2026-07-20 23:08:11 +00:00
committed by GitHub
co-authored by haosdent Claude Opus 4.7 Justin Wood
parent b7c20d0cfa
commit 58b2012aa2
4 changed files with 73 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/language/pooling/test_reward.py
+13 -2
View File
@@ -9,6 +9,7 @@ import torch.nn.functional as F
from transformers import AutoModel
from vllm.platforms import current_platform
from vllm.utils.mem_constants import MiB_bytes
from ....conftest import HfRunner
from ....utils import VLLM_PATH
@@ -106,7 +107,12 @@ def test_prm_models(
if current_platform.is_cpu():
pytest.skip("CPU only supports V1")
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
with vllm_runner(
model,
max_model_len=1024,
dtype=dtype,
kv_cache_memory_bytes=64 * MiB_bytes,
) as vllm_model:
vllm_outputs = vllm_model.token_classify(math_step_prompts)
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
@@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs(
if not FIXTURE_REWARD_RESULT.get(model):
pytest.skip(f"No available golden outputs for {model}.")
with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
with vllm_runner(
model,
max_model_len=1024,
dtype=dtype,
kv_cache_memory_bytes=64 * MiB_bytes,
) as vllm_model:
vllm_outputs = vllm_model.token_classify(math_step_prompts)
golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
tests/utils_/test_mem_utils.py
+17
View File
@@ -67,6 +67,23 @@ def test_memory_profiling():
non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa
assert abs(non_torch_ratio - 1) <= 0.05
assert result.torch_peak_increase == 1024 * 1024 * 1024
expected_total_consumed = (256 + 512) * 1024 * 1024
total_consumed_ratio = result.total_consumed / expected_total_consumed
assert abs(total_consumed_ratio - 1) <= 0.05, (
f"total_consumed={result.total_consumed}, "
f"expected={expected_total_consumed}, "
f"ratio={total_consumed_ratio}"
)
expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
assert abs(non_kv_ratio - 1) <= 0.05, (
f"non_kv_cache_memory={result.non_kv_cache_memory}, "
f"expected={expected_non_kv}, "
f"ratio={non_kv_ratio}"
)
del weights
lib.cudaFree(handle1)
lib.cudaFree(handle2)
vllm/utils/mem_utils.py
+22 -9
View File
@@ -110,6 +110,7 @@ class MemorySnapshot:
"""Memory snapshot."""
torch_peak: int = 0
torch_allocated: int = 0
free_memory: int = 0
total_memory: int = 0
cuda_memory: int = 0
@@ -139,9 +140,9 @@ class MemorySnapshot:
# After `torch.accelerator.reset_peak_memory_stats()`,
# `torch.accelerator.memory_reserved()` will keep growing, and only shrink
# when we call `torch.accelerator.empty_cache()` or OOM happens.
self.torch_peak = torch.accelerator.memory_stats(device).get(
"allocated_bytes.all.peak", 0
)
stats = torch.accelerator.memory_stats(device)
self.torch_peak = stats.get("allocated_bytes.all.peak", 0)
self.torch_allocated = stats.get("allocated_bytes.all.current", 0)
self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device)
if current_platform.is_integrated_gpu(device.index):
@@ -172,6 +173,7 @@ class MemorySnapshot:
return MemorySnapshot(
torch_peak=self.torch_peak - other.torch_peak,
torch_allocated=self.torch_allocated - other.torch_allocated,
free_memory=self.free_memory - other.free_memory,
total_memory=self.total_memory - other.total_memory,
cuda_memory=self.cuda_memory - other.cuda_memory,
@@ -185,6 +187,7 @@ class MemorySnapshot:
def __repr__(self) -> str:
return (
f"torch_peak={format_gib(self.torch_peak)}GiB, "
f"torch_allocated={format_gib(self.torch_allocated)}GiB, "
f"free_memory={format_gib(self.free_memory)}GiB, "
f"total_memory={format_gib(self.total_memory)}GiB, "
f"{current_platform.device_name}_memory={format_gib(self.cuda_memory)}GiB, "
@@ -202,6 +205,8 @@ class MemoryProfilingResult:
non_kv_cache_memory: int = 0
torch_peak_increase: int = 0
non_torch_increase: int = 0
total_consumed: int = 0
transient_peak_headroom: int = 0
weights_memory: int = 0
before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
profile_time: float = 0.0
@@ -219,8 +224,8 @@ class MemoryProfilingResult:
f"{format_gib(self.non_kv_cache_memory)}GiB; "
f"torch peak memory increase: "
f"{format_gib(self.torch_peak_increase)}GiB; "
f"non-torch forward increase memory: "
f"{format_gib(self.non_torch_increase)}GiB; "
f"total consumed (from mem_get_info): "
f"{format_gib(self.total_consumed)}GiB; "
f"weights memory: {format_gib(self.weights_memory)}GiB."
)
@@ -306,8 +311,16 @@ def memory_profiling(
result.non_torch_increase = diff_from_create.non_torch_memory
result.profile_time = diff_profile.timestamp
non_torch_memory = result.non_torch_increase
peak_activation_memory = result.torch_peak_increase
result.non_kv_cache_memory = (
non_torch_memory + peak_activation_memory + result.weights_memory
# Measure total consumption via mem_get_info() instead of
# memory_reserved(), which goes negative when pluggable allocators
# (e.g. cumem) bypass PyTorch's tracking.
result.total_consumed = (
result.before_create.free_memory - result.after_profile.free_memory
)
# total_consumed already covers persistent torch allocations; add only the
# transient peak headroom to avoid double-counting.
result.transient_peak_headroom = (
result.after_profile.torch_peak - result.after_profile.torch_allocated
)
result.non_kv_cache_memory = result.total_consumed + result.transient_peak_headroom
vllm/v1/worker/gpu_worker.py
+21 -35
View File
@@ -488,33 +488,18 @@ class Worker(WorkerBase):
) as profile_result:
self.model_runner.profile_run()
profile_torch_peak = torch.accelerator.memory_stats(self.device).get(
"allocated_bytes.all.peak", 0
)
# Profile CUDA graph memory if graphs will be captured.
# ROCm is included: #44825 moved the profiler to
# torch.accelerator.get_memory_info (reliable on ROCm, as used by
# the AMD-CI mem tests), and graph_pool_handle resolves to the same
# torch.cuda handle the live capture path already uses on ROCm.
# XPU stays excluded (see #39977).
cudagraph_memory_estimate = 0
if (
current_platform.is_cuda_alike()
and self.vllm_config.compilation_config.cudagraph_mode
!= CUDAGraphMode.NONE
):
cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()
# Use the pre-cudagraph torch peak to avoid double-counting.
profile_result.torch_peak_increase = (
profile_torch_peak - profile_result.before_profile.torch_peak
)
profile_result.non_kv_cache_memory = (
profile_result.non_torch_increase
+ profile_result.torch_peak_increase
+ profile_result.weights_memory
)
# Profile CUDA graph memory if graphs will be captured.
# ROCm is included: #44825 moved the profiler to
# torch.accelerator.get_memory_info (reliable on ROCm, as used by
# the AMD-CI mem tests), and graph_pool_handle resolves to the same
# torch.cuda handle the live capture path already uses on ROCm.
# XPU stays excluded (see #39977).
cudagraph_memory_estimate = 0
if (
current_platform.is_cuda_alike()
and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
):
cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()
# Respect the opt-in flag as originally designed.
cudagraph_memory_estimate_applied = (
@@ -523,8 +508,10 @@ class Worker(WorkerBase):
else 0
)
self.non_torch_memory = profile_result.non_torch_increase
self.peak_activation_memory = profile_result.torch_peak_increase
self.total_consumed = profile_result.total_consumed
self.peak_activation_memory = (
profile_result.transient_peak_headroom + cudagraph_memory_estimate_applied
)
self.cudagraph_memory_estimate = cudagraph_memory_estimate
free_gpu_memory = profile_result.after_profile.free_memory
@@ -817,9 +804,8 @@ class Worker(WorkerBase):
redundancy_buffer_memory = 150 * (1 << 20)
non_kv_cache_memory = (
self.model_runner.model_memory_usage
self.total_consumed
+ self.peak_activation_memory
+ self.non_torch_memory
+ cuda_graph_memory_bytes
)
kv_cache_memory_bytes_to_gpu_limit = (
@@ -840,10 +826,10 @@ class Worker(WorkerBase):
f"Desired GPU memory utilization is "
f"({self.cache_config.gpu_memory_utilization}, "
f"{format_gib(self.requested_memory)} GiB). "
f"Actual usage is {format_gib(self.model_runner.model_memory_usage)} "
f"GiB for weight, {format_gib(self.peak_activation_memory)} GiB "
f"for peak activation, {format_gib(self.non_torch_memory)} GiB "
f"for non-torch memory, and {format_gib(cuda_graph_memory_bytes)} "
f"Actual usage is {format_gib(self.total_consumed)} "
f"GiB for consumed memory (weights + non-torch), "
f"{format_gib(self.peak_activation_memory)} GiB "
f"for peak activation, and {format_gib(cuda_graph_memory_bytes)} "
f"GiB for CUDAGraph memory. Replace gpu_memory_utilization "
f"config with `--kv-cache-memory="
f"{kv_cache_memory_bytes_to_requested_limit}` "