[Bugfix] Fix pooled Whisper sliding-window KV sizing (#47071)
Signed-off-by: Andy Lo <andy@mistral.ai>
Andy Lo committed
Jul 1, 2026 at 10:33 UTC
fa4bec90567715d1482b01fb6ad6c16bc5c1ed36
2 files changed
+43
-2
tests/models/multimodal/generation/test_voxtral_realtime.py
+32
-1
@@ -15,11 +15,14 @@ from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
15
from vllm import LLM, EngineArgs, SamplingParams
16
from vllm.assets.audio import AudioAsset
17
from vllm.engine.arg_utils import AsyncEngineArgs
18
+from vllm.utils.math_utils import cdiv
19
from vllm.v1.engine.async_llm import AsyncLLM
20
+from vllm.v1.kv_cache_interface import SlidingWindowSpec
21
22
from ....utils import ROCM_ENGINE_KWARGS
23
24
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
25
+AUDIO_LAYER_NAME = "whisper_encoder.whisper_encoder.layers.0.layers.self_attn.attn"
26
ENGINE_CONFIG = {
27
"model": MODEL_NAME,
28
"max_model_len": 8192,
@@ -60,6 +63,31 @@ def _normalize(texts: list[str]) -> list[str]:
63
return texts
64
65
66
+def assert_encoder_kv_cache_spec(engine: LLM) -> None:
67
+ vllm_config = engine.llm_engine.vllm_config
68
+ audio_config = vllm_config.model_config.hf_config.audio_config
69
+ kv_cache_specs_per_rank = engine.llm_engine.model_executor.get_kv_cache_specs()
70
+
71
+ assert len(kv_cache_specs_per_rank) == 1
72
+ kv_cache_specs = kv_cache_specs_per_rank[0]
73
+ assert AUDIO_LAYER_NAME in kv_cache_specs, kv_cache_specs.keys()
74
+ spec = kv_cache_specs[AUDIO_LAYER_NAME]
75
+
76
+ assert audio_config.sliding_window == 750
77
+ assert audio_config.block_pool_size == 4
78
+ assert isinstance(spec, SlidingWindowSpec)
79
+ assert spec.block_size == 16
80
+ assert spec.num_kv_heads == 128
81
+ assert spec.sliding_window == cdiv(750, 4) == 188
82
+ assert (
83
+ spec.max_admission_blocks_per_request(
84
+ max_num_batched_tokens=1,
85
+ max_model_len=vllm_config.model_config.max_model_len,
86
+ )
87
+ == 13
88
+ )
89
+
90
+
91
@pytest.fixture
92
def audio_assets() -> list[AudioAsset]:
93
return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
@@ -71,7 +99,9 @@ def tokenizer() -> MistralTokenizer:
99
100
101
@pytest.fixture
74
-def engine():
102
+def engine(monkeypatch: pytest.MonkeyPatch):
103
+ # Disable multiprocessing allows us to access model executor from LLM engine
104
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
105
engine_args = EngineArgs(**ENGINE_CONFIG)
106
llm = LLM.from_engine_args(engine_args)
107
try:
@@ -95,6 +125,7 @@ async def async_engine():
125
126
127
def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
128
+ assert_encoder_kv_cache_spec(engine)
129
audio_config = tokenizer.instruct_tokenizer.tokenizer.audio
130
131
def from_file(file_path: str):
vllm/model_executor/models/whisper_causal.py
+11
-1
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
23
from vllm.model_executor.layers.rotary_embedding import get_rope
24
from vllm.model_executor.models.mistral import MistralMLP
25
from vllm.model_executor.models.whisper import WhisperPosEmbedType
26
+from vllm.utils.math_utils import cdiv
27
from vllm.v1.attention.backend import (
28
AttentionBackend,
29
AttentionMetadata,
@@ -39,7 +40,7 @@ except ImportError:
40
from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
41
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
42
from vllm.v1.attention.selector import get_attn_backend
42
-from vllm.v1.kv_cache_interface import AttentionSpec
43
+from vllm.v1.kv_cache_interface import AttentionSpec, SlidingWindowSpec
44
45
from .utils import make_layers
46
@@ -329,6 +330,15 @@ class WhisperCausalAttentionWithBlockPooling(Attention):
330
kv_cache_spec,
331
num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
332
)
333
+ if isinstance(kv_cache_spec, SlidingWindowSpec):
334
+ # The KV cache manager counts blocks in pooled units, so express the
335
+ # window in pooled units too to avoid reserving `block_pool_size`x
336
+ # too many blocks. The kernel window is unaffected because it comes
337
+ # from the attention impl, not this spec.
338
+ kv_cache_spec = replace(
339
+ kv_cache_spec,
340
+ sliding_window=cdiv(kv_cache_spec.sliding_window, self.block_pool_size),
341
+ )
342
return kv_cache_spec
343
344