Skip to content

fix: #448 drops inference mode during vLLM memory profiling #460

Description

@shipiyouniao

Problem

#448 replaces vLLM's GPUWorker.determine_available_memory() but does not
preserve the original method's @torch.inference_mode() decorator.

In vLLM 0.22.1, GPUModelRunner.profile_run() is not independently decorated
with @torch.inference_mode(); it relies on the outer worker method to establish
inference mode. As a result, the patched path runs the multimodal dummy profile
with autograd enabled, retains the vision encoder's backward graph and
intermediate activations, and can turn startup into a real CUDA OOM.

This is a regression introduced by #448 rather than an unavoidable multimodal
profiling peak.

Evidence

An A/B test used the same image, model, startup arguments and profiling inputs.
Only kvcached/integration/vllm/patches.py changed.

Both paths received the same inputs:

  • V1 model runner
  • pixel_values=(65536, 1536), bfloat16
  • image_grid_thw=(1, 3)
  • identical multimodal budgets and item counts

Memory observed around the same model.embed_multimodal() call:

Path Before embed After embed PyTorch peak Process footprint
Before #448 0.867 GiB 0.944 GiB 1.943 GiB about 2.9 GiB
With #448 0.867 GiB 13.13 GiB 13.62 GiB about 14.06 GiB

The increase occurs inside the vision encoder forward, before
determine_available_memory() returns and before KV cache configuration,
num_blocks, or _allocate_kv_cache_tensors() runs. This rules out eager KV
allocation, Mamba cache allocation, changed dummy inputs, and image resizing as
the cause.

The relevant execution contract is:

# vLLM 0.22.1
@torch.inference_mode()
def determine_available_memory(self) -> int:
    ...
    self.model_runner.profile_run()

The replacement introduced by #448 calls the same profiling methods without
that decorator. The expected mode difference is therefore:

  • original path: torch.is_grad_enabled() == False and
    torch.is_inference_mode_enabled() == True
  • patched path: torch.is_grad_enabled() == True and
    torch.is_inference_mode_enabled() == False

Proposed fix

Preserve vLLM's original execution contract by decorating the replacement with
@torch.inference_mode() (or entering the equivalent context) for the complete
profiling path.

This must not skip the model forward, multimodal profiling, CUDA graph
profiling/capture, or real CUDA error propagation. It only prevents autograd
state from being retained during an inference-only startup profile.

Regression tests should assert that both profile_run() and
profile_cudagraph_memory() execute with inference mode enabled when called
through the patched determine_available_memory().

Expected behavior

The patched worker should retain #448's logical KV-capacity behavior while
matching vLLM's original inference-mode semantics and avoiding the additional
vision-encoder activation peak.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions