[copy of #45208] CuMem slept-L1 fragmentation accounting (#49208)
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>
Matthew Bonanni committed
Jul 20, 2026 at 19:08 UTC
58b2012aa26f2b85560ecd6988c8fa2a773804c1
4 files changed
+73
-46
tests/models/language/pooling/test_reward.py
+13
-2
@@ -9,6 +9,7 @@ import torch.nn.functional as F
9
from transformers import AutoModel
10
11
from vllm.platforms import current_platform
12
+from vllm.utils.mem_constants import MiB_bytes
13
14
from ....conftest import HfRunner
15
from ....utils import VLLM_PATH
@@ -106,7 +107,12 @@ def test_prm_models(
107
if current_platform.is_cpu():
108
pytest.skip("CPU only supports V1")
109
109
- with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
110
+ with vllm_runner(
111
+ model,
112
+ max_model_len=1024,
113
+ dtype=dtype,
114
+ kv_cache_memory_bytes=64 * MiB_bytes,
115
+ ) as vllm_model:
116
vllm_outputs = vllm_model.token_classify(math_step_prompts)
117
118
with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
@@ -145,7 +151,12 @@ def test_prm_models_with_golden_outputs(
151
if not FIXTURE_REWARD_RESULT.get(model):
152
pytest.skip(f"No available golden outputs for {model}.")
153
148
- with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
154
+ with vllm_runner(
155
+ model,
156
+ max_model_len=1024,
157
+ dtype=dtype,
158
+ kv_cache_memory_bytes=64 * MiB_bytes,
159
+ ) as vllm_model:
160
vllm_outputs = vllm_model.token_classify(math_step_prompts)
161
162
golden_outputs = load_reward_outputs(FIXTURE_REWARD_RESULT[model])
tests/utils_/test_mem_utils.py
+17
@@ -67,6 +67,23 @@ def test_memory_profiling():
67
non_torch_ratio = result.non_torch_increase / (256 * 1024 * 1024) # noqa
68
assert abs(non_torch_ratio - 1) <= 0.05
69
assert result.torch_peak_increase == 1024 * 1024 * 1024
70
+
71
+ expected_total_consumed = (256 + 512) * 1024 * 1024
72
+ total_consumed_ratio = result.total_consumed / expected_total_consumed
73
+ assert abs(total_consumed_ratio - 1) <= 0.05, (
74
+ f"total_consumed={result.total_consumed}, "
75
+ f"expected={expected_total_consumed}, "
76
+ f"ratio={total_consumed_ratio}"
77
+ )
78
+
79
+ expected_non_kv = expected_total_consumed + 1024 * 1024 * 1024
80
+ non_kv_ratio = result.non_kv_cache_memory / expected_non_kv
81
+ assert abs(non_kv_ratio - 1) <= 0.05, (
82
+ f"non_kv_cache_memory={result.non_kv_cache_memory}, "
83
+ f"expected={expected_non_kv}, "
84
+ f"ratio={non_kv_ratio}"
85
+ )
86
+
87
del weights
88
lib.cudaFree(handle1)
89
lib.cudaFree(handle2)
vllm/utils/mem_utils.py
+22
-9
@@ -110,6 +110,7 @@ class MemorySnapshot:
110
"""Memory snapshot."""
111
112
torch_peak: int = 0
113
+ torch_allocated: int = 0
114
free_memory: int = 0
115
total_memory: int = 0
116
cuda_memory: int = 0
@@ -139,9 +140,9 @@ class MemorySnapshot:
140
# After `torch.accelerator.reset_peak_memory_stats()`,
141
# `torch.accelerator.memory_reserved()` will keep growing, and only shrink
142
# when we call `torch.accelerator.empty_cache()` or OOM happens.
142
- self.torch_peak = torch.accelerator.memory_stats(device).get(
143
- "allocated_bytes.all.peak", 0
144
- )
143
+ stats = torch.accelerator.memory_stats(device)
144
+ self.torch_peak = stats.get("allocated_bytes.all.peak", 0)
145
+ self.torch_allocated = stats.get("allocated_bytes.all.current", 0)
146
147
self.free_memory, self.total_memory = torch.accelerator.get_memory_info(device)
148
if current_platform.is_integrated_gpu(device.index):
@@ -172,6 +173,7 @@ class MemorySnapshot:
173
174
return MemorySnapshot(
175
torch_peak=self.torch_peak - other.torch_peak,
176
+ torch_allocated=self.torch_allocated - other.torch_allocated,
177
free_memory=self.free_memory - other.free_memory,
178
total_memory=self.total_memory - other.total_memory,
179
cuda_memory=self.cuda_memory - other.cuda_memory,
@@ -185,6 +187,7 @@ class MemorySnapshot:
187
def __repr__(self) -> str:
188
return (
189
f"torch_peak={format_gib(self.torch_peak)}GiB, "
190
+ f"torch_allocated={format_gib(self.torch_allocated)}GiB, "
191
f"free_memory={format_gib(self.free_memory)}GiB, "
192
f"total_memory={format_gib(self.total_memory)}GiB, "
193
f"{current_platform.device_name}_memory={format_gib(self.cuda_memory)}GiB, "
@@ -202,6 +205,8 @@ class MemoryProfilingResult:
205
non_kv_cache_memory: int = 0
206
torch_peak_increase: int = 0
207
non_torch_increase: int = 0
208
+ total_consumed: int = 0
209
+ transient_peak_headroom: int = 0
210
weights_memory: int = 0
211
before_create: MemorySnapshot = field(default_factory=MemorySnapshot)
212
profile_time: float = 0.0
@@ -219,8 +224,8 @@ class MemoryProfilingResult:
224
f"{format_gib(self.non_kv_cache_memory)}GiB; "
225
f"torch peak memory increase: "
226
f"{format_gib(self.torch_peak_increase)}GiB; "
222
- f"non-torch forward increase memory: "
223
- f"{format_gib(self.non_torch_increase)}GiB; "
227
+ f"total consumed (from mem_get_info): "
228
+ f"{format_gib(self.total_consumed)}GiB; "
229
f"weights memory: {format_gib(self.weights_memory)}GiB."
230
)
231
@@ -306,8 +311,16 @@ def memory_profiling(
311
result.non_torch_increase = diff_from_create.non_torch_memory
312
result.profile_time = diff_profile.timestamp
313
309
- non_torch_memory = result.non_torch_increase
310
- peak_activation_memory = result.torch_peak_increase
311
- result.non_kv_cache_memory = (
312
- non_torch_memory + peak_activation_memory + result.weights_memory
314
+ # Measure total consumption via mem_get_info() instead of
315
+ # memory_reserved(), which goes negative when pluggable allocators
316
+ # (e.g. cumem) bypass PyTorch's tracking.
317
+ result.total_consumed = (
318
+ result.before_create.free_memory - result.after_profile.free_memory
319
+ )
320
+
321
+ # total_consumed already covers persistent torch allocations; add only the
322
+ # transient peak headroom to avoid double-counting.
323
+ result.transient_peak_headroom = (
324
+ result.after_profile.torch_peak - result.after_profile.torch_allocated
325
)
326
+ result.non_kv_cache_memory = result.total_consumed + result.transient_peak_headroom
vllm/v1/worker/gpu_worker.py
+21
-35
@@ -488,33 +488,18 @@ class Worker(WorkerBase):
488
) as profile_result:
489
self.model_runner.profile_run()
490
491
- profile_torch_peak = torch.accelerator.memory_stats(self.device).get(
492
- "allocated_bytes.all.peak", 0
493
- )
494
-
495
- # Profile CUDA graph memory if graphs will be captured.
496
- # ROCm is included: #44825 moved the profiler to
497
- # torch.accelerator.get_memory_info (reliable on ROCm, as used by
498
- # the AMD-CI mem tests), and graph_pool_handle resolves to the same
499
- # torch.cuda handle the live capture path already uses on ROCm.
500
- # XPU stays excluded (see #39977).
501
- cudagraph_memory_estimate = 0
502
- if (
503
- current_platform.is_cuda_alike()
504
- and self.vllm_config.compilation_config.cudagraph_mode
505
- != CUDAGraphMode.NONE
506
- ):
507
- cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()
508
-
509
- # Use the pre-cudagraph torch peak to avoid double-counting.
510
- profile_result.torch_peak_increase = (
511
- profile_torch_peak - profile_result.before_profile.torch_peak
512
- )
513
- profile_result.non_kv_cache_memory = (
514
- profile_result.non_torch_increase
515
- + profile_result.torch_peak_increase
516
- + profile_result.weights_memory
517
- )
491
+ # Profile CUDA graph memory if graphs will be captured.
492
+ # ROCm is included: #44825 moved the profiler to
493
+ # torch.accelerator.get_memory_info (reliable on ROCm, as used by
494
+ # the AMD-CI mem tests), and graph_pool_handle resolves to the same
495
+ # torch.cuda handle the live capture path already uses on ROCm.
496
+ # XPU stays excluded (see #39977).
497
+ cudagraph_memory_estimate = 0
498
+ if (
499
+ current_platform.is_cuda_alike()
500
+ and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
501
+ ):
502
+ cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()
503
504
# Respect the opt-in flag as originally designed.
505
cudagraph_memory_estimate_applied = (
@@ -523,8 +508,10 @@ class Worker(WorkerBase):
508
else 0
509
)
510
526
- self.non_torch_memory = profile_result.non_torch_increase
527
- self.peak_activation_memory = profile_result.torch_peak_increase
511
+ self.total_consumed = profile_result.total_consumed
512
+ self.peak_activation_memory = (
513
+ profile_result.transient_peak_headroom + cudagraph_memory_estimate_applied
514
+ )
515
self.cudagraph_memory_estimate = cudagraph_memory_estimate
516
517
free_gpu_memory = profile_result.after_profile.free_memory
@@ -817,9 +804,8 @@ class Worker(WorkerBase):
804
redundancy_buffer_memory = 150 * (1 << 20)
805
806
non_kv_cache_memory = (
820
- self.model_runner.model_memory_usage
807
+ self.total_consumed
808
+ self.peak_activation_memory
822
- + self.non_torch_memory
809
+ cuda_graph_memory_bytes
810
)
811
kv_cache_memory_bytes_to_gpu_limit = (
@@ -840,10 +826,10 @@ class Worker(WorkerBase):
826
f"Desired GPU memory utilization is "
827
f"({self.cache_config.gpu_memory_utilization}, "
828
f"{format_gib(self.requested_memory)} GiB). "
843
- f"Actual usage is {format_gib(self.model_runner.model_memory_usage)} "
844
- f"GiB for weight, {format_gib(self.peak_activation_memory)} GiB "
845
- f"for peak activation, {format_gib(self.non_torch_memory)} GiB "
846
- f"for non-torch memory, and {format_gib(cuda_graph_memory_bytes)} "
829
+ f"Actual usage is {format_gib(self.total_consumed)} "
830
+ f"GiB for consumed memory (weights + non-torch), "
831
+ f"{format_gib(self.peak_activation_memory)} GiB "
832
+ f"for peak activation, and {format_gib(cuda_graph_memory_bytes)} "
833
f"GiB for CUDAGraph memory. Replace gpu_memory_utilization "
834
f"config with `--kv-cache-memory="
835
f"{kv_cache_memory_bytes_to_requested_limit}` "