[Hybrid]: Decouple Logical Block Size from Physical Page Size #24486

zhiyuan1i · 2025-09-09T06:52:08Z

Purpose

This PR introduces a hybrid cache architecture that separates logical kernel block size from
physical page size, enabling more flexible memory management. Key changes include:

Added kernel_block_size field to CacheConfig for logical block sizing
Enhanced platform-specific configurations for CUDA and ROCm to support hybrid blocks
Implemented block table conversion logic between physical and logical representations
Added support for different physical/logical block size ratios in V1 worker components

This hybrid model decoupling enables independent development of high-performance operators
without being constrained by linear attention mechanisms like Mamba, addressing performance
bottlenecks discussed in issues #24280 and
#23161.

Test Plan

Added comprehensive tests in tests/v1/worker/test_gpu_model_runner.py to verify:

Block table conversion between physical and logical representations
Proper handling of different block size ratios
Integration with existing GPU model runner functionality
Platform-specific configurations for CUDA and ROCm

Test Result

pytest tests/v1/worker/test_gpu_model_runner.py - 20 passes

tests/v1/worker/test_gpu_model_runner.py ....................                                                                                                                        [100%]

===================================================================================== warnings summary =====================================================================================
../../../../opt/conda/envs/vllm-upstream/lib/python3.12/site-packages/torch/cuda/__init__.py:63
  /opt/conda/envs/vllm-upstream/lib/python3.12/site-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
    import pynvml  # type: ignore[import]

<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

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

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================================================================== 20 passed, 3 warnings in 89.20s (0:01:29) =========================================================================
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 introduces a hybrid cache architecture to decouple logical and physical block sizes, which is a significant enhancement for memory management. The changes span configuration, platform-specific code, and the core block table management. The implementation in block_table.py appears solid. However, I've identified some critical issues in the tests intended to validate this new functionality. The tests are flawed and do not correctly verify the hybrid block logic, which could mask bugs. Additionally, there's a piece of logic in the GPUModelRunner that could be made more robust. My review focuses on fixing these test and implementation issues to ensure the new feature is reliable and well-tested.

tests/v1/worker/test_gpu_model_runner.py

vllm/v1/worker/gpu_model_runner.py

heheda12345 · 2025-09-09T07:30:44Z

Also CC @tdoublep

heheda12345

Discussed with @zhiyuan1i offline. Two major concerns:

I prefer to calculate kernel block size for each attention backend in gpu_model_runner
would be great if BlockTable.block_table and BlockTable.physical_block_table can be merged into one tensor.

zhiyuan1i · 2025-09-09T09:31:44Z

@heheda12345 Thanks for the prompt feedback! I’ve addressed suggestion2 and merged BlockTable.block_table and BlockTable.physical_block_table into a single tensor as recommended. :)

vllm/platforms/cuda.py

vllm/v1/worker/gpu_model_runner.py

vllm/v1/worker/block_table.py

tjtanaa · 2025-09-11T03:52:24Z

CC @gshtras @hongxiayang as this also affect ROCm

vllm/v1/worker/block_table.py

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

mergify · 2025-09-25T03:05:29Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhiyuan1i.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

vllm/platforms/cuda.py

vllm/v1/worker/gpu_model_runner.py

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

heheda12345 · 2025-09-30T18:57:38Z

vllm/model_executor/models/config.py

-        # with mamba layers, use FlashInfer instead).
-        attn_block_size = 16 * cdiv(mamba_page_size,
-                                    16 * attn_page_size_1_token)
+        # Model may be marked as is_hybrid


which model?

heheda12345 · 2025-09-30T19:01:02Z

vllm/model_executor/models/config.py

        # override attention block size if either (a) the
        # user has not set it or (b) the user has set it
        # too small.
+        #


nit: revert?

heheda12345 · 2025-09-30T19:01:57Z

vllm/platforms/cuda.py

        # Note: model_config may be None during testing
-        if model_config is not None and model_config.use_mla:
+        # Note: block_size is initialized in
+        # HybridAttentionMambaModelConfig.verify_and_update_config


Suggested change

# HybridAttentionMambaModelConfig.verify_and_update_config

# HybridAttentionMambaModelConfig.verify_and_update_config

# for models with both attention and mamba,

heheda12345 · 2025-09-30T19:09:38Z

vllm/v1/worker/block_table.py

            # Use a "virtual block" which equals to world_size * block_size
            # for block_table_indices calculation.
            virtual_block_size = self.block_size * self.dcp_world_size
+


can you revert?

heheda12345 · 2025-09-30T21:49:24Z

vllm/v1/worker/gpu_model_runner.py

                    has_attn = True
+                    kv_manager_block_size = kv_cache_spec.block_size
+                    logical_kernel_size_list = \
+                        self._find_compatible_block_sizes(


Why not _select_common_block_size? And prefer to select the value once and pass it to both self.may_reinitialize_input_batch and _reshape_kv_cache_tensors to avoid silent bugs.

heheda12345 · 2025-09-30T21:51:57Z

vllm/v1/worker/gpu_model_runner.py

                if isinstance(kv_cache_spec, AttentionSpec):
                    has_attn = True
+                    kv_manager_block_size = kv_cache_spec.block_size
+                    logical_kernel_size_list = \


please avoid using "logical". Is "kernel_block_size" OK in this case?

heheda12345 · 2025-09-30T21:53:36Z

vllm/v1/worker/gpu_model_runner.py

+                        self._find_compatible_block_sizes(
+                        kv_manager_block_size, attn_backend, return_all=False)
+                    logical_kernel_size = logical_kernel_size_list[0]
+                    num_blocks_per_phys_block = (kv_manager_block_size //


please avoid "physical block"

mergify · 2025-09-30T21:54:12Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @zhiyuan1i.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

zhiyuan1i requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac, alexm-redhat, simon-mo, youkaichao, mgoin, tlrmchlsmth, houseroad, hmellor, yewentao256 and ProExpertProg as code owners September 9, 2025 06:52

mergify bot added rocm Related to AMD ROCm v1 labels Sep 9, 2025

gemini-code-assist bot reviewed Sep 9, 2025

View reviewed changes

tests/v1/worker/test_gpu_model_runner.py Outdated Show resolved Hide resolved

tests/v1/worker/test_gpu_model_runner.py Outdated Show resolved Hide resolved

vllm/v1/worker/gpu_model_runner.py Outdated Show resolved Hide resolved

zhiyuan1i force-pushed the hybrid-cache-groups branch from 4e3eeca to 8f2ee3d Compare September 9, 2025 06:58

heheda12345 reviewed Sep 9, 2025

View reviewed changes

zhiyuan1i force-pushed the hybrid-cache-groups branch from 954ade4 to 0e0823a Compare September 9, 2025 09:28

zhiyuan1i force-pushed the hybrid-cache-groups branch 2 times, most recently from 6d1735e to 0b544bf Compare September 9, 2025 14:43

heheda12345 reviewed Sep 10, 2025

View reviewed changes

vllm/platforms/cuda.py Outdated Show resolved Hide resolved

vllm/v1/worker/gpu_model_runner.py Outdated Show resolved Hide resolved

vllm/v1/worker/block_table.py Show resolved Hide resolved

vllm/v1/worker/block_table.py Outdated Show resolved Hide resolved

zhiyuan1i force-pushed the hybrid-cache-groups branch from 62e5072 to 55e2235 Compare September 10, 2025 07:36

heheda12345 reviewed Sep 11, 2025

View reviewed changes

zhiyuan1i force-pushed the hybrid-cache-groups branch from 55e2235 to b0e1d3b Compare September 11, 2025 06:25

zhiyuan1i requested a review from zhuohan123 as a code owner September 11, 2025 06:25

heheda12345 reviewed Sep 23, 2025

View reviewed changes

vllm/v1/worker/block_table.py Outdated Show resolved Hide resolved

zhiyuan1i added 3 commits September 24, 2025 04:01

minor fixs

adba4a5

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

fix kernel block size for attn groups

279d1d0

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

Merge remote-tracking branch 'upstream/main' into hybrid-cache-groups

5d328d2

mergify bot removed the needs-rebase label Sep 24, 2025

heheda12345 added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 24, 2025

zhiyuan1i added 3 commits September 24, 2025 04:24

fix lint

7cb4fc3

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

fix test

585f2bf

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

rename

413272b

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

zhiyuan1i requested a review from ApostaC as a code owner September 24, 2025 12:45

zhiyuan1i force-pushed the hybrid-cache-groups branch from 16a6dc0 to 771ef36 Compare September 25, 2025 03:04

mergify bot added the needs-rebase label Sep 25, 2025

fix mla issue

a51673b

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

zhiyuan1i force-pushed the hybrid-cache-groups branch from 771ef36 to a51673b Compare September 25, 2025 03:25

Merge remote-tracking branch 'upstream/main' into hybrid-cache-groups

5691f12

mergify bot removed the needs-rebase label Sep 25, 2025

heheda12345 reviewed Sep 26, 2025

View reviewed changes

vllm/platforms/cuda.py Outdated Show resolved Hide resolved

heheda12345 reviewed Sep 26, 2025

View reviewed changes

vllm/v1/worker/gpu_model_runner.py Show resolved Hide resolved

zhiyuan1i added 4 commits September 26, 2025 09:53

fix flashinfer mla

d865f00

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

try to use default block_size

dd7bfc8

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

Merge branch 'main' into hybrid-cache-groups

6a97abb

Merge branch 'main' into hybrid-cache-groups

74e0ff1

zhiyuan1i requested a review from LucasWilkinson as a code owner September 27, 2025 08:23

zhiyuan1i added 2 commits September 27, 2025 11:52

fix bugs in mla

c9231e8

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

fix corner case

248dbd5

Signed-off-by: lizhiyuan <uniartisan2017@gmail.com>

zhiyuan1i force-pushed the hybrid-cache-groups branch from f07921a to 248dbd5 Compare September 29, 2025 06:48

heheda12345 reviewed Sep 30, 2025

View reviewed changes

mergify bot added the needs-rebase label Sep 30, 2025

	# HybridAttentionMambaModelConfig.verify_and_update_config
	# HybridAttentionMambaModelConfig.verify_and_update_config
	# for models with both attention and mamba,

Uh oh!

[Hybrid]: Decouple Logical Block Size from Physical Page Size #24486

Are you sure you want to change the base?

[Hybrid]: Decouple Logical Block Size from Physical Page Size #24486

Conversation

zhiyuan1i commented Sep 9, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

heheda12345 commented Sep 9, 2025

Uh oh!

heheda12345 left a comment

Choose a reason for hiding this comment

Uh oh!

zhiyuan1i commented Sep 9, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tjtanaa commented Sep 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Sep 25, 2025

Uh oh!

Uh oh!

Uh oh!

heheda12345 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

heheda12345 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

heheda12345 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

heheda12345 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

heheda12345 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

heheda12345 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

heheda12345 Sep 30, 2025

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Sep 30, 2025

Uh oh!

Uh oh!

zhiyuan1i commented Sep 9, 2025 •

edited by github-actions bot

Loading

tjtanaa commented Sep 11, 2025 •

edited

Loading