[Bugfix] Fix pooled Whisper encoder sliding-window kernel size (#47437)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 2, 2026 at 18:19 UTC
443e68cfa616312d88aecd613a8013d5b5169956
2 files changed
+27
-13
tests/models/multimodal/generation/test_voxtral_realtime.py
+3
-1
@@ -78,7 +78,9 @@ def assert_encoder_kv_cache_spec(engine: LLM) -> None:
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
81
+ # cdiv(750, 4) == 188 pooled tokens cover the model's window; the extra
82
+ # +1 is an eviction margin (see whisper_causal.py get_kv_cache_spec).
83
+ assert spec.sliding_window == cdiv(750, 4) + 1 == 189
84
assert (
85
spec.max_admission_blocks_per_request(
86
max_num_batched_tokens=1,
vllm/model_executor/models/whisper_causal.py
+24
-12
@@ -112,7 +112,9 @@ class WhisperCausalConv1d(nn.Conv1d):
112
113
@functools.lru_cache
114
def create_whisper_attention_backend_with_block_pooling(
115
- underlying_attn_backend: AttentionBackend, block_pool_size: int
115
+ underlying_attn_backend: AttentionBackend,
116
+ block_pool_size: int,
117
+ sliding_window: int | None = None,
118
) -> type[AttentionBackend]:
119
prefix = "WhisperCausalAttentionWithBlockPooling_"
120
underlying_builder = underlying_attn_backend.get_builder_cls()
@@ -127,11 +129,22 @@ def create_whisper_attention_backend_with_block_pooling(
129
device: torch.device,
130
):
131
assert kv_cache_spec.num_kv_heads % block_pool_size == 0
130
- kv_cache_spec = replace(
131
- kv_cache_spec,
132
+ # Scale pooled-unit quantities back to unpooled (encoder-token)
133
+ # units for the underlying attention metadata: the KV cache spec
134
+ # counts blocks in pooled units, but the kernel runs on the
135
+ # `block_pool_size`x-expanded sequence built below.
136
+ spec_overrides = dict(
137
block_size=kv_cache_spec.block_size * block_pool_size,
138
num_kv_heads=kv_cache_spec.num_kv_heads // block_pool_size,
139
)
140
+ if isinstance(kv_cache_spec, SlidingWindowSpec):
141
+ # The manager keeps `sliding_window` in pooled units; the kernel
142
+ # runs on the expanded sequence, so give it the model's window in
143
+ # unpooled units (`get_kv_cache_spec` sizes the pooled window so
144
+ # this window always stays resident).
145
+ assert sliding_window is not None
146
+ spec_overrides["sliding_window"] = sliding_window
147
+ kv_cache_spec = replace(kv_cache_spec, **spec_overrides)
148
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
149
# Override model_config-derived values with the actual
150
# encoder values from kv_cache_spec
@@ -303,7 +316,7 @@ class WhisperCausalAttentionWithBlockPooling(Attention):
316
attn_type=attn_type,
317
)
318
attn_backend = create_whisper_attention_backend_with_block_pooling(
306
- underlying_attn_backend, block_pool_size
319
+ underlying_attn_backend, block_pool_size, per_layer_sliding_window
320
)
321
322
super().__init__(
@@ -331,14 +344,13 @@ class WhisperCausalAttentionWithBlockPooling(Attention):
344
num_kv_heads=self.block_pool_size * kv_cache_spec.num_kv_heads,
345
)
346
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
- )
347
+ # The manager counts blocks in pooled units, so express the window in
348
+ # pooled units to avoid reserving `block_pool_size`x too many blocks.
349
+ # The `+1` is an eviction margin: block-aligned eviction only keeps
350
+ # `pooled - 1` pooled tokens, so this guarantees the kernel's full
351
+ # (unpooled) window still stays resident.
352
+ pooled = cdiv(kv_cache_spec.sliding_window, self.block_pool_size) + 1
353
+ kv_cache_spec = replace(kv_cache_spec, sliding_window=pooled)
354
return kv_cache_spec
355
356