[BugFix] Correct max memory usage for multiple KV-cache groups

[peakcrosser7]

Purpose

This PR fixes a calculation error in _max_memory_usage_bytes_from_groups() that led to underestimating total memory usage in multi-group cases.

Root Cause

The original implementation only calculated blocks_needed based on the first KV-Cache group (index 0). In models with multiple KV-Cache groups (e.g., hybrid Mamba models), the block usage from other groups was neglected. This resulted in a lower-than-actual total memory estimate, also caused the engine to over-calculate the max_model_len.

For example, in test_auto_fit_max_model_len_with_hybrid(), the model consists of both Mamba and FullAttn groups. Since the original logic only considered the Mamba group (index 0), which occupies a fixed number of blocks when prefix-caching is disabled, it ignored the memory usage of the FullAttn group and led to an excessively large max_model_len.

Solution

Updated the logic to sum the blocks from all KV-Cache groups, ensuring the total memory usage accurately reflects the requirements of the entire model.

Test Plan

Added a new unit test test_auto_fit_max_model_len_with_hybrid()

pytest tests/v1/core/test_kv_cache_utils.py::test_auto_fit_max_model_len_with_hybrid

Test Result

Before fix:

_____________________________________________________ ERROR collecting tests/v1/core/test_kv_cache_utils.py ______________________________________________________
tests/v1/core/test_kv_cache_utils.py:2127: in <module>
    test_auto_fit_max_model_len_with_hybrid()
tests/v1/core/test_kv_cache_utils.py:2038: in test_auto_fit_max_model_len_with_hybrid
    assert vllm_config.model_config.max_model_len == 1024
E   AssertionError: assert 8192 == 1024
E    +  where 8192 = ModelConfig(model='Qwen/Qwen3-0.6B', model_weights='', runner='auto', convert='auto', tokenizer='Qwen/Qwen3-0.6B', tok...ide_attention_dtype=None, logits_processors=None, io_processor_plugin=None, pooler_config=None, multimodal_config=None).max_model_len
E    +    where ModelConfig(model='Qwen/Qwen3-0.6B', model_weights='', runner='auto', convert='auto', tokenizer='Qwen/Qwen3-0.6B', tok...ide_attention_dtype=None, logits_processors=None, io_processor_plugin=None, pooler_config=None, multimodal_config=None) = VllmConfig(model_config=ModelConfig(model='Qwen/Qwen3-0.6B', model_weights='', runner='auto', convert='auto', tokenize...6572044341470', optimization_level=<OptimizationLevel.O2: 2>, performance_mode='balanced', weight_transfer_config=None).model_config

After fix:

root@hhy_develop_vllm_opsrc:~/huanghy/vllm_opsrc# pytest tests/v1/core/test_kv_cache_utils.py::test_auto_fit_max_model_len_with_hybrid
====================================================================== test session starts =======================================================================
platform linux -- Python 3.11.13, pytest-8.4.2, pluggy-1.6.0
rootdir: /root/huanghy/vllm_opsrc
configfile: pyproject.toml
plugins: hypothesis-6.137.1, anyio-4.11.0
collected 1 item                                                                                                                                                 

tests/v1/core/test_kv_cache_utils.py .                                                                                                                     [100%]

======================================================================== warnings summary ========================================================================
<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:241
  <frozen importlib._bootstrap>:241: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

../../../opt/conda/lib/python3.11/site-packages/torch/jit/_script.py:362: 14 warnings
  /opt/conda/lib/python3.11/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================================================================ 1 passed, 16 warnings in 13.92s =================================================================
sys:1: DeprecationWarning: builtin type swigvarlink has no __module__ attribute

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request addresses a bug in the maximum memory usage calculation, which is particularly relevant for hybrid models. The change in vllm/v1/core/kv_cache_utils.py correctly sums the memory requirements across all KV cache groups, whereas the previous implementation only considered the first group. This ensures accurate memory estimation, for instance when auto-fitting max_model_len. The addition of a new test case for hybrid KV cache specs in tests/v1/core/test_kv_cache_utils.py is a good measure to prevent regressions.

[heheda12345]

LGTM!

[repne]

Hi @peakcrosser7, thank you for your work on this. I've been testing this PR without issues for the past week or so.
I was wondering, is this PR superseeded by #37124?

[swtb3]

Hi @peakcrosser7, thank you for your work on this. I've been testing this PR without issues for the past week or so. I was wondering, is this PR superseeded by #37124?

Hello, im author of #37124 , I also fixed _max_memory_usage_bytes_from_groups as part of a broader set of changes to hybrid Mamba/attention KV cache handling (reporting, allocation, and concurrency estimation). There is overlap on that one function. Let me know how we should proceed :)

(my apologies as i did not see this PR as the issue that brought me here was to do with memory over estimation for the qwen 3.5 series.)

[peakcrosser7]

Hi @repne, thanks for pointing that out! It seems there is indeed some overlap between this PR and #37124.

Hi @swtb3, great work on #37124! It looks like a more comprehensive solution to the problem. I'm not entirely sure about the best way to move forward with my PR, since the auto-merge wasn’t triggered as expected earlier.

@heheda12345, what are your thoughts on this?

[repne]

From a quick test I am getting OOM with #37124, but not with #36030. So the highlighted lines seems to be needed @swtb3. This on Blackwell + Qwen3.5-27B-FP8

[swtb3]

Ive done some digging on the cause of the OOM. I think that to get the proper allocation for Qwen3.5 will have me back to the drawing board. It may not be as simple as I first thought. I would say, if this PR is ready and tested then go for it. I will rebase my PR on top and continue figuring it out. If youve any thoughts on the OOM lets discuss over on #37124

[repne]

Will use this PR for the time being, ping me when you have something you want me to test. Thank you!

[swtb3]

@repne new changes pushed, could you test? cheers!

[tdoublep]

LGTM