[Bug]: v1 fails when token budget is set to 128K

[rajveerb]

Your current environment

The output of `python collect_env.py`

INFO 03-25 03:03:43 [__init__.py:256] Automatically detected platform cuda.
Collecting environment information...
PyTorch version: 2.6.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.31.3
Libc version: glibc-2.31

Python version: 3.10.16 (main, Dec 11 2024, 16:24:50) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-204-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A40
GPU 1: NVIDIA A40

Nvidia driver version: 525.85.12
cuDNN version: Probably one of the following:
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.6.0
/usr/local/cuda-11.8/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.6.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn.so.8.6.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.6.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.6.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.6.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.6.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.6.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.6.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
Address sizes:                      48 bits physical, 48 bits virtual
CPU(s):                             128
On-line CPU(s) list:                0-127
Thread(s) per core:                 2
Core(s) per socket:                 32
Socket(s):                          2
NUMA node(s):                       2
Vendor ID:                          AuthenticAMD
CPU family:                         25
Model:                              1
Model name:                         AMD EPYC 7513 32-Core Processor
Stepping:                           1
Frequency boost:                    enabled
CPU MHz:                            1768.982
CPU max MHz:                        2600.0000
CPU min MHz:                        1500.0000
BogoMIPS:                           5200.31
Virtualization:                     AMD-V
L1d cache:                          2 MiB
L1i cache:                          2 MiB
L2 cache:                           32 MiB
L3 cache:                           256 MiB
NUMA node0 CPU(s):                  0-31,64-95
NUMA node1 CPU(s):                  32-63,96-127
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-cusparselt-cu12==0.6.2
[pip3] nvidia-ml-py==12.570.86
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] pyzmq==26.2.1
[pip3] torch==2.6.0
[pip3] torchaudio==2.6.0
[pip3] torchvision==0.21.0
[pip3] transformers==4.48.3
[pip3] triton==3.2.0
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-cusparselt-cu12    0.6.2                    pypi_0    pypi
[conda] nvidia-ml-py              12.570.86                pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] pyzmq                     26.2.1                   pypi_0    pypi
[conda] torch                     2.6.0                    pypi_0    pypi
[conda] torchaudio                2.6.0                    pypi_0    pypi
[conda] torchvision               0.21.0                   pypi_0    pypi
[conda] transformers              4.48.3                   pypi_0    pypi
[conda] triton                    3.2.0                    pypi_0    pypi
ROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: 0.8.2.dev0+g61c7a1b85.d20250325
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
GPU0    GPU1    CPU Affinity    NUMA Affinity
GPU0     X      NODE    32-63,96-127    1
GPU1    NODE     X      32-63,96-127    1

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks

LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64
NCCL_P2P_LEVEL=NVL
NCCL_CUMEM_ENABLE=0
TORCHINDUCTOR_COMPILE_THREADS=1
CUDA_MODULE_LOADING=LAZY

🐛 Describe the bug

In v1 is enabled in v0.8.1, the program runs into OOM when max_num_batched_tokens (i.e. max token budget for the scheduler) is set to 128K for Llama 3.1 8B.

There is enough KV space based on below log:

INFO 03-25 03:01:40 [kv_cache_utils.py:540] Maximum concurrency for 131,072 tokens per request: 3.05x

Here's a minimal python script to replicate the bug:

# command: python test.py
import asyncio, time, uuid, pytest
from vllm.engine.arg_utils import EngineArgs
from vllm.platforms import current_platform
from vllm.usage.usage_lib import UsageContext
from vllm.v1.engine.core_client import EngineCoreClient, AsyncMPClient
from vllm.v1.executor.abstract import Executor
from typing import List
from vllm.v1.engine import EngineCoreRequest
from vllm import SamplingParams

if not current_platform.is_cuda():
    pytest.skip(reason="V1 currently only supported on CUDA.",
                allow_module_level=True)

MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"


def create_dummy_token_ids(size: int, vocab_size: int) -> List[int]:
    """Create a list of dummy token IDs of specified size.
    
    Uses a simple pattern that cycles through tokens 0 to vocab_size-1 to ensure
    we have valid token IDs within the model's vocabulary range while being
    deterministic.
    
    Args:
        size: Number of token IDs to generate
    vocab_size: Size of the model's vocabulary

    Returns:
        List of token IDs within the valid vocabulary range
    """
    return [i % vocab_size for i in range(size)]  # Cycle through 0 to vocab_size-1
        
def create_dummy_request(token_count: int, vocab_size: int) -> EngineCoreRequest:
    """Create a dummy EngineCoreRequest with specified number of tokens.

    Args:
        token_count: Number of tokens in the request
        vocab_size: Size of the model's vocabulary
    """
    return EngineCoreRequest(
        request_id=str(uuid.uuid4()),
        prompt=None,  # Not needed when added to EngineCoreClient
        prompt_token_ids=create_dummy_token_ids(token_count, vocab_size),
        mm_inputs=None,
        mm_hashes=None,
        mm_placeholders=None,
        sampling_params=SamplingParams(max_tokens=1),
        eos_token_id=None,
        arrival_time=time.time(),
        lora_request=None,
    )

@pytest.mark.asyncio
async def test_engine_core_client_asyncio(monkeypatch):

    with monkeypatch.context() as m:
        m.setenv("VLLM_USE_V1", "1")
        # Set max_num_batched_tokens to max model length
        engine_args = EngineArgs(model=MODEL_NAME,
				enable_chunked_prefill=False,
                max_num_batched_tokens=131072,
				gpu_memory_utilization=0.95,
                tensor_parallel_size=2,
                enable_prefix_caching=False,
			)
        vllm_config = engine_args.create_engine_config(
            usage_context=UsageContext.UNKNOWN_CONTEXT)
        executor_class = Executor.get_class(vllm_config)
        client: AsyncMPClient = EngineCoreClient.make_client(
            multiprocess_mode=True,
            asyncio_mode=True,
            vllm_config=vllm_config,
            executor_class=executor_class,
            log_stats=True
        )

        await client.add_request_async(create_dummy_request(131072, vllm_config.model_config.get_vocab_size()))

        # loop over get_output_async until request is complete
        all_finished = False
        while not all_finished:
            output = await client.get_output_async()
            if output is None:
                continue
            
            print(output)
            # Check if all requests are finished
            all_finished = all(out.finished for out in output.outputs)

        client.shutdown()


if __name__ == "__main__":
    monkeypatch = pytest.MonkeyPatch()
    asyncio.run(test_engine_core_client_asyncio(monkeypatch))

OOM message:

ERROR 03-25 03:02:03 [core.py:425]     raise result
ERROR 03-25 03:02:03 [core.py:425] torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.75 GiB. GPU 0 has a total capacity of 44.40 GiB of which 541.56 MiB is free. Including non-PyTorch memory, this process has 43.87 GiB memory in use. Of the allocated memory 40.08 GiB is allocated by PyTorch, with 73.00 MiB allocated in private pools (e.g., CUDA Graphs), and 1.01 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
ERROR 03-25 03:02:03 [core.py:425] 
CRITICAL 03-25 03:02:03 [core_client.py:269] Got fatal signal from worker processes, shutting down. See stack trace above for root cause issue.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[Qubitium]

@rajveerb Reduce gpu_memory_utilization=0.95 to 0.90 and try agin.

[rajveerb]

@Qubitium Thanks for pointing that.

gpu_memory_utilization=0.9 didn't work, if I send one request with 128K and after it finishes, I send another 128K length request. It leads to OOM despite KV cache usage being 0 after the first request finishes, it's quite bizarre.

Setting gpu_memory_utilization=0.85 did resolve it though.

But, I don't understand why 0.9 didn't work.

[Qubitium]

@rajveerb

But, I don't understand why 0.9 didn't work.

I run into these oom myself from time to time and I believe it likely has to do with memory fragmentation. vllm usually does a good job of estimating vram since it knows the max context length, dtype, kv dtype, and precalculate kv size, max batch size based on all info but one thing it can't predict is contiguous memory. Some parts of the code/kernel requires contiguous memory and some are fragmented (v1 prefix caching). So gpu may show 40GB used but maybe a portion of that is wasted due to fragmentation. I mean, gpu has to allocate memory just like linux kernel does and no matter how smart it becomes, there is waste. This is my guess.

Here is the ref showing that V1 prefix caching is fragmenting memory (likely on purpose): #15319

[rajveerb]

@Qubitium

Thanks for sharing.

I assumed that profile run has high fidelity and would ideally take care of this issue.

Surprising to me since there was almost 2x space available for 128K context length requests,

[Qubitium]

I assumed that profile run has high fidelity and would ideally take care of this issue.

Surprising to me since there was almost 2x space available for 128K context length requests,

@rajveerb Maybe the existing profile run on startup is not hitting and testing the prefix cache and hitting chunk prefill paths? Worth looking into

[coopslarhette]

also running into this running gemma 3 27b:

INFO 03-27 18:56:51 [kv_cache_utils.py:566] GPU KV cache size: 153,200 tokens
INFO 03-27 18:56:51 [kv_cache_utils.py:569] Maximum concurrency for 131,072 tokens per request: 1.17x

will try dropping gpu_memory_utilization, but thought it was interesting

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[github-actions]

This issue has been automatically closed due to inactivity. Please feel free to reopen if you feel it is still relevant. Thank you!