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, 13:30 UTC
[Bugfix][Spec Decode] Fix extract_hidden_states for VLM models (#38987)
Signed-off-by: Aaron Batilo <abatilo@coreweave.com>
Aaron Batilo committed
Apr 5, 2026 at 03:41 UTC
9a528260ef648500262709550807c292098a70c0
2 files changed
+173
tests/v1/spec_decode/test_extract_hidden_states.py
+163
@@ -1,11 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import json
5
from unittest import mock
6
7
import numpy as np
8
import pytest
9
import torch
10
+from transformers import CLIPVisionConfig, LlamaConfig, LlavaConfig, PretrainedConfig
11
12
from tests.v1.attention.utils import (
13
BatchSpec,
@@ -23,6 +25,10 @@ from vllm.config import (
25
)
26
from vllm.config.load import LoadConfig
27
from vllm.platforms import current_platform
28
+from vllm.transformers_utils.config import get_hf_text_config
29
+from vllm.transformers_utils.configs.extract_hidden_states import (
30
+ ExtractHiddenStatesConfig,
31
+)
32
from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
33
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
34
@@ -323,3 +329,160 @@ def test_propose_different_layer_counts(num_hidden_layers):
329
330
assert draft_tokens.shape == (batch_size, 1)
331
assert torch.equal(draft_tokens, sampled_token_ids)
332
+
333
+
334
+# ---------------------------------------------------------------------------
335
+# VLM / composite config tests for ExtractHiddenStatesConfig
336
+# ---------------------------------------------------------------------------
337
+
338
+
339
+class _DummyVLMConfig(PretrainedConfig):
340
+ """Minimal composite config that mimics VLMs like Kimi-K2.5 or LLaVA.
341
+
342
+ The text model's parameters (hidden_size, num_attention_heads, …) live
343
+ exclusively under ``text_config``; the top-level config has none of them.
344
+ """
345
+
346
+ model_type = "test_vlm"
347
+
348
+ def __init__(self, text_config: PretrainedConfig, **kwargs):
349
+ self.text_config = text_config
350
+ super().__init__(architectures=["LlamaForCausalLM"], **kwargs)
351
+
352
+ def get_text_config(self, decoder: bool = False) -> PretrainedConfig:
353
+ del decoder
354
+ return self.text_config
355
+
356
+
357
+def test_extract_hidden_states_text_only_config_regression():
358
+ """Text-only models (no nested text_config) must keep working."""
359
+ model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
360
+
361
+ speculative_config = SpeculativeConfig(
362
+ target_model_config=model_config,
363
+ target_parallel_config=ParallelConfig(),
364
+ method="extract_hidden_states",
365
+ num_speculative_tokens=1,
366
+ draft_model_config={
367
+ "hf_config": {
368
+ "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4],
369
+ }
370
+ },
371
+ )
372
+
373
+ assert speculative_config.draft_model_config is not None
374
+ # For text-only models, hf_text_config should be the config itself.
375
+ assert speculative_config.draft_model_config.hf_text_config is (
376
+ speculative_config.draft_model_config.hf_config
377
+ )
378
+ assert (
379
+ speculative_config.draft_model_config.hf_text_config.num_attention_heads
380
+ == model_config.hf_text_config.num_attention_heads
381
+ )
382
+
383
+
384
+def test_extract_hidden_states_config_preserves_vlm_text_config():
385
+ """A real VLM config (LLaVA) with nested text_config must be preserved."""
386
+ text_config = LlamaConfig(
387
+ vocab_size=32000,
388
+ hidden_size=128,
389
+ intermediate_size=256,
390
+ num_hidden_layers=2,
391
+ num_attention_heads=8,
392
+ )
393
+ vlm_config = LlavaConfig(
394
+ vision_config=CLIPVisionConfig(),
395
+ text_config=text_config,
396
+ )
397
+
398
+ # Precondition: to_dict() flattens the nested config to a plain dict.
399
+ assert isinstance(vlm_config.to_dict()["text_config"], dict)
400
+
401
+ extract_config = ExtractHiddenStatesConfig(
402
+ vlm_config,
403
+ eagle_aux_hidden_state_layer_ids=[1, 2],
404
+ )
405
+
406
+ # The fix: text_config is still a PretrainedConfig, not a dict.
407
+ assert isinstance(extract_config.text_config, LlamaConfig)
408
+
409
+ extracted = get_hf_text_config(extract_config)
410
+ assert extracted is extract_config.text_config
411
+ assert extracted.num_attention_heads == text_config.num_attention_heads
412
+ assert extracted.hidden_size == text_config.hidden_size
413
+
414
+ # Serialization must still round-trip correctly.
415
+ serialized = extract_config.to_dict()
416
+ assert isinstance(serialized["text_config"], dict)
417
+ assert serialized["text_config"]["num_attention_heads"] == (
418
+ text_config.num_attention_heads
419
+ )
420
+
421
+ json_str = json.loads(extract_config.to_json_string())
422
+ assert json_str["text_config"]["num_attention_heads"] == (
423
+ text_config.num_attention_heads
424
+ )
425
+
426
+
427
+def test_extract_hidden_states_speculative_config_vlm():
428
+ """SpeculativeConfig with a VLM target must build without errors."""
429
+ nested_text_config = LlamaConfig(
430
+ vocab_size=32000,
431
+ hidden_size=128,
432
+ intermediate_size=256,
433
+ num_hidden_layers=2,
434
+ num_attention_heads=8,
435
+ )
436
+
437
+ target_model_config = ModelConfig(
438
+ model=model_dir,
439
+ runner="generate",
440
+ max_model_len=100,
441
+ )
442
+ # Replace the real text-only config with our composite VLM config.
443
+ target_model_config.hf_config = _DummyVLMConfig(
444
+ text_config=nested_text_config,
445
+ )
446
+ target_model_config.hf_text_config = nested_text_config
447
+
448
+ speculative_config = SpeculativeConfig(
449
+ target_model_config=target_model_config,
450
+ target_parallel_config=ParallelConfig(),
451
+ method="extract_hidden_states",
452
+ num_speculative_tokens=1,
453
+ draft_model_config={
454
+ "hf_config": {
455
+ "eagle_aux_hidden_state_layer_ids": [1, 2],
456
+ }
457
+ },
458
+ )
459
+
460
+ assert speculative_config.draft_model_config is not None
461
+ assert isinstance(
462
+ speculative_config.draft_model_config.hf_config.text_config,
463
+ LlamaConfig,
464
+ )
465
+ assert speculative_config.draft_model_config.hf_text_config is (
466
+ speculative_config.draft_model_config.hf_config.text_config
467
+ )
468
+ assert (
469
+ speculative_config.draft_model_config.hf_text_config.num_attention_heads
470
+ == nested_text_config.num_attention_heads
471
+ )
472
+
473
+
474
+def test_extract_hidden_states_config_invalid_text_config():
475
+ """A nested text_config missing required attrs must still be rejected."""
476
+ broken_text_config = PretrainedConfig(hidden_size=128)
477
+ vlm_config = _DummyVLMConfig(text_config=broken_text_config)
478
+
479
+ extract_config = ExtractHiddenStatesConfig(
480
+ vlm_config,
481
+ eagle_aux_hidden_state_layer_ids=[1],
482
+ )
483
+
484
+ # The object is preserved (not flattened), …
485
+ assert extract_config.text_config is broken_text_config
486
+ # … but validation still rejects the missing attribute.
487
+ with pytest.raises(ValueError, match="num_attention_heads"):
488
+ get_hf_text_config(extract_config)
vllm/transformers_utils/configs/extract_hidden_states.py
+10
@@ -23,10 +23,14 @@ class ExtractHiddenStatesConfig(PretrainedConfig):
23
24
if isinstance(model, dict):
25
model_dict = model
26
+ source_text_config = None
27
elif isinstance(model, PretrainedConfig):
28
model_dict = model.to_dict()
29
+ text_config = model.get_text_config()
30
+ source_text_config = text_config if text_config is not model else None
31
else:
32
model_dict = {}
33
+ source_text_config = None
34
35
# Combine: model_dict first, then kwargs override
36
combined = {**model_dict, **kwargs}
@@ -35,6 +39,12 @@ class ExtractHiddenStatesConfig(PretrainedConfig):
39
40
combined["architectures"] = ["ExtractHiddenStatesModel"]
41
42
+ # to_dict() and kwargs both flatten text_config to a plain dict;
43
+ # downstream get_hf_text_config() needs it as a PretrainedConfig
44
+ # for attribute access. Re-insert the original object.
45
+ if source_text_config is not None:
46
+ combined["text_config"] = source_text_config
47
+
48
super().__init__(**combined)
49
50
@classmethod