A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 19:45 UTC
Fix ambiguous num_blocks for hybrid attn mamba (#37236)
Signed-off-by: Collin McCarthy <cmccarthy@nvidia.com> Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>
Collin McCarthy committed
Mar 30, 2026 at 04:09 UTC
1031c84c360874ddc37589f21b15686788ac142e
2 files changed
+51
-14
tests/v1/worker/test_gpu_model_runner.py
+31
-1
@@ -1,6 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from types import SimpleNamespace
5
+
6
import numpy as np
7
import pytest
8
import torch
@@ -30,6 +32,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum
32
from vllm.v1.core.kv_cache_utils import estimate_max_model_len, get_kv_cache_configs
33
from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
34
from vllm.v1.kv_cache_interface import (
35
+ AttentionSpec,
36
FullAttentionSpec,
37
KVCacheConfig,
38
KVCacheGroupSpec,
@@ -38,7 +41,7 @@ from vllm.v1.kv_cache_interface import (
41
from vllm.v1.sample.metadata import SamplingMetadata
42
from vllm.v1.worker.gpu_input_batch import InputBatch
43
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
41
-from vllm.v1.worker.utils import select_common_block_size
44
+from vllm.v1.worker.utils import AttentionGroup, select_common_block_size
45
46
BLOCK_SIZE = 16
47
NUM_BLOCKS = 10
@@ -946,6 +949,33 @@ def test_hybrid_attention_mamba_tensor_shapes():
949
assert torch.equal(actual_ssm, expected_ssm)
950
951
952
+def test_update_hybrid_attention_mamba_layout_with_num_block_2_rewrites_stride():
953
+ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
954
+
955
+ ambiguous_cache = torch.empty((2, 2, BLOCK_SIZE, 1, 8), dtype=torch.float16)
956
+ """Ambiguous, because both dims[0=kv_dim] and dims[1=num_blocks] == 2"""
957
+ hidden_size = ambiguous_cache.shape[2:].numel()
958
+ assert ambiguous_cache.stride()[:2] == (2 * hidden_size, hidden_size)
959
+
960
+ attention_spec = AttentionSpec(
961
+ block_size=BLOCK_SIZE, num_kv_heads=1, head_size=8, dtype=torch.float16
962
+ )
963
+ runner_stub = SimpleNamespace(
964
+ cache_config=SimpleNamespace(cache_dtype="auto"),
965
+ _kv_cache_spec_attn_group_iterator=lambda: iter(
966
+ [AttentionGroup(FlashAttentionBackend, ["attn"], attention_spec, 0)]
967
+ ),
968
+ )
969
+ GPUModelRunner._update_hybrid_attention_mamba_layout(
970
+ runner_stub, {"attn": ambiguous_cache}, [BLOCK_SIZE]
971
+ )
972
+
973
+ assert ambiguous_cache.stride()[:2] == (hidden_size, 2 * hidden_size), """\
974
+ We expect _update_hybrid_attention_mamba_layout to re-stride the cache from:
975
+ (2, num_blocks) -> (num_blocks, 2), even when num_blocks==2,
976
+ which was ambiguous before get_kv_cache_block_dim was used"""
977
+
978
+
979
def test_hybrid_block_table_initialization():
980
"""Test hybrid block table with different kernel and kvcache_manager block
981
sizes."""
vllm/v1/worker/gpu_model_runner.py
+20
-13
@@ -6653,12 +6653,12 @@ class GPUModelRunner(
6653
raise NotImplementedError
6654
6655
if has_attn and has_mamba:
6656
- self._update_hybrid_attention_mamba_layout(kv_caches)
6656
+ self._update_hybrid_attention_mamba_layout(kv_caches, kernel_block_sizes)
6657
6658
return kv_caches
6659
6660
def _update_hybrid_attention_mamba_layout(
6661
- self, kv_caches: dict[str, torch.Tensor]
6661
+ self, kv_caches: dict[str, torch.Tensor], kernel_block_sizes: list[int]
6662
) -> None:
6663
"""
6664
Update the layout of attention layers from (2, num_blocks, ...) to
@@ -6666,23 +6666,30 @@ class GPUModelRunner(
6666
6667
Args:
6668
kv_caches: The KV cache buffer of each layer.
6669
+ kernel_block_sizes: The kernel block sizes for each KV cache group.
6670
"""
6671
6672
for group in self._kv_cache_spec_attn_group_iterator():
6673
kv_cache_spec = group.kv_cache_spec
6674
+ if not isinstance(kv_cache_spec, AttentionSpec):
6675
+ continue
6676
+ block_dim = group.backend.get_kv_cache_block_dim(
6677
+ kernel_block_sizes[group.kv_cache_group_id],
6678
+ kv_cache_spec.num_kv_heads,
6679
+ kv_cache_spec.head_size,
6680
+ cache_dtype_str=self.cache_config.cache_dtype,
6681
+ )
6682
+ # block_dim: 0 means (num_blocks, 2, ...); 1 means (2, num_blocks, ...).
6683
+ if block_dim == 0:
6684
+ continue
6685
+ assert block_dim == 1
6686
for layer_name in group.layer_names:
6687
kv_cache = kv_caches[layer_name]
6675
- if isinstance(kv_cache_spec, AttentionSpec) and kv_cache.shape[0] == 2:
6676
- assert kv_cache.shape[1] != 2, (
6677
- "Fail to determine whether the layout is "
6678
- "(2, num_blocks, ...) or (num_blocks, 2, ...) for "
6679
- f"a tensor of shape {kv_cache.shape}"
6680
- )
6681
- hidden_size = kv_cache.shape[2:].numel()
6682
- kv_cache.as_strided_(
6683
- size=kv_cache.shape,
6684
- stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]),
6685
- )
6688
+ hidden_size = kv_cache.shape[2:].numel()
6689
+ kv_cache.as_strided_(
6690
+ size=kv_cache.shape,
6691
+ stride=(hidden_size, 2 * hidden_size, *kv_cache.stride()[2:]),
6692
+ )
6693
6694
def initialize_kv_cache_tensors(
6695
self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]