| 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, |
| 14 | create_common_attn_metadata, |
| 15 | ) |
| 16 | from vllm.config import ( |
| 17 | AttentionConfig, |
| 18 | CacheConfig, |
| 19 | DeviceConfig, |
| 20 | ModelConfig, |
| 21 | ParallelConfig, |
| 22 | SchedulerConfig, |
| 23 | SpeculativeConfig, |
| 24 | VllmConfig, |
| 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 | |
| 35 | model_dir = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" |
| 36 | DEVICE_TYPE = current_platform.device_type |
| 37 | |
| 38 | |
| 39 | def _create_proposer( |
| 40 | num_speculative_tokens: int = 1, |
| 41 | layer_ids: list[int] | None = None, |
| 42 | ) -> ExtractHiddenStatesProposer: |
| 43 | """Create an ExtractHiddenStatesProposer for testing.""" |
| 44 | if layer_ids is None: |
| 45 | layer_ids = [1, 2, 3, 4] |
| 46 | |
| 47 | model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100) |
| 48 | |
| 49 | speculative_config = SpeculativeConfig( |
| 50 | target_model_config=model_config, |
| 51 | target_parallel_config=ParallelConfig(), |
| 52 | method="extract_hidden_states", |
| 53 | num_speculative_tokens=num_speculative_tokens, |
| 54 | draft_model_config={ |
| 55 | "hf_config": { |
| 56 | "eagle_aux_hidden_state_layer_ids": layer_ids, |
| 57 | } |
| 58 | }, |
| 59 | ) |
| 60 | |
| 61 | device = DEVICE_TYPE |
| 62 | vllm_config = VllmConfig( |
| 63 | model_config=model_config, |
| 64 | cache_config=CacheConfig(), |
| 65 | speculative_config=speculative_config, |
| 66 | device_config=DeviceConfig(device=device), |
| 67 | parallel_config=ParallelConfig(), |
| 68 | load_config=LoadConfig(), |
| 69 | scheduler_config=SchedulerConfig( |
| 70 | max_model_len=model_config.max_model_len, |
| 71 | is_encoder_decoder=model_config.is_encoder_decoder, |
| 72 | ), |
| 73 | attention_config=AttentionConfig(), |
| 74 | ) |
| 75 | |
| 76 | return ExtractHiddenStatesProposer(vllm_config=vllm_config, device=device) |
| 77 | |
| 78 | |
| 79 | def test_proposer_initialization(): |
| 80 | """Test that the proposer initializes correctly with the right parameters.""" |
| 81 | layer_ids = [1, 2, 3, 4] |
| 82 | proposer = _create_proposer(num_speculative_tokens=1, layer_ids=layer_ids) |
| 83 | |
| 84 | assert proposer.num_hidden_states == len(layer_ids) |
| 85 | assert proposer.vllm_config.speculative_config is not None |
| 86 | assert proposer.vllm_config.speculative_config.num_speculative_tokens == 1 |
| 87 | |
| 88 | # Verify the hidden states buffer is correctly shaped |
| 89 | expected_shape = ( |
| 90 | proposer.max_num_tokens, |
| 91 | len(layer_ids), |
| 92 | proposer.hidden_size, |
| 93 | ) |
| 94 | assert proposer.hidden_states.shape == expected_shape |
| 95 | |
| 96 | |
| 97 | def test_proposer_initialization_missing_layer_ids(): |
| 98 | """Test that initialization fails when layer_ids are not provided.""" |
| 99 | model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100) |
| 100 | |
| 101 | speculative_config = SpeculativeConfig( |
| 102 | target_model_config=model_config, |
| 103 | target_parallel_config=ParallelConfig(), |
| 104 | method="extract_hidden_states", |
| 105 | num_speculative_tokens=1, |
| 106 | draft_model_config={ |
| 107 | "hf_config": {} # Missing eagle_aux_hidden_state_layer_ids |
| 108 | }, |
| 109 | ) |
| 110 | |
| 111 | device = DEVICE_TYPE |
| 112 | vllm_config = VllmConfig( |
| 113 | model_config=model_config, |
| 114 | cache_config=CacheConfig(), |
| 115 | speculative_config=speculative_config, |
| 116 | device_config=DeviceConfig(device=device), |
| 117 | parallel_config=ParallelConfig(), |
| 118 | load_config=LoadConfig(), |
| 119 | scheduler_config=SchedulerConfig( |
| 120 | max_model_len=model_config.max_model_len, |
| 121 | is_encoder_decoder=model_config.is_encoder_decoder, |
| 122 | ), |
| 123 | attention_config=AttentionConfig(), |
| 124 | ) |
| 125 | |
| 126 | with pytest.raises( |
| 127 | ValueError, match="eagle_aux_hidden_state_layer_ids must be set" |
| 128 | ): |
| 129 | ExtractHiddenStatesProposer(vllm_config=vllm_config, device=device) |
| 130 | |
| 131 | |
| 132 | def test_prepare_next_token_ids_padded(): |
| 133 | """Test for prepare_next_token_ids_padded with extract_hidden_states. |
| 134 | |
| 135 | Since num_speculative_tokens == 1, sampled_token_ids has shape (batch_size, 1). |
| 136 | For each request we either use the sampled token (if valid and not discarded) |
| 137 | or a backup token from the request state. |
| 138 | """ |
| 139 | device = torch.device(DEVICE_TYPE) |
| 140 | |
| 141 | num_requests = 4 |
| 142 | req_ids = [f"req_{i + 1}" for i in range(num_requests)] |
| 143 | mock_input_batch = mock.MagicMock(spec=InputBatch) |
| 144 | mock_input_batch.req_ids = req_ids |
| 145 | mock_input_batch.num_reqs = num_requests |
| 146 | mock_input_batch.vocab_size = 100 |
| 147 | mock_input_batch.num_tokens_no_spec = np.array([5] * num_requests) |
| 148 | |
| 149 | mock_requests: dict[str, CachedRequestState] = {} |
| 150 | for req_id in req_ids: |
| 151 | mock_request = mock.MagicMock(spec=CachedRequestState) |
| 152 | # Each request will have a backup next token id of 10, 20, 30, 40 |
| 153 | mock_request.get_token_id.return_value = int(req_id.split("_")[1]) * 10 |
| 154 | mock_requests[req_id] = mock_request |
| 155 | |
| 156 | # explicitly discard the last request |
| 157 | discarded_req_mask = torch.tensor( |
| 158 | [False, False, False, True], dtype=torch.bool, device=device |
| 159 | ) |
| 160 | |
| 161 | # With num_speculative_tokens=1, sampled_token_ids has shape [batch_size, 1] |
| 162 | sampled_token_ids = torch.tensor( |
| 163 | [ |
| 164 | [1], # valid, use 1 |
| 165 | [4], # valid, use 4 |
| 166 | [-1], # invalid, use backup token "30" |
| 167 | [2], # explicitly discarded, use backup token "40" |
| 168 | ], |
| 169 | dtype=torch.int32, |
| 170 | device=device, |
| 171 | ) |
| 172 | |
| 173 | expected_next_token_ids_cpu = [1, 4, 30, 40] |
| 174 | expected_next_token_ids_tensor = torch.tensor( |
| 175 | expected_next_token_ids_cpu, dtype=torch.int32, device=device |
| 176 | ) |
| 177 | |
| 178 | proposer = _create_proposer(num_speculative_tokens=1) |
| 179 | |
| 180 | # valid_sampled_tokens_count tracks if token is valid (not -1 and in vocab range) |
| 181 | # It doesn't depend on whether the request is discarded |
| 182 | expected_valid_sampled_tokens_count = torch.tensor( |
| 183 | [1, 1, 0, 1], dtype=torch.int32, device=device |
| 184 | ) |
| 185 | |
| 186 | next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded( |
| 187 | sampled_token_ids, |
| 188 | mock_requests, |
| 189 | mock_input_batch, |
| 190 | discarded_req_mask, |
| 191 | ) |
| 192 | |
| 193 | assert torch.equal(next_token_ids, expected_next_token_ids_tensor) |
| 194 | assert torch.equal(valid_sampled_tokens_count, expected_valid_sampled_tokens_count) |
| 195 | |
| 196 | |
| 197 | def test_propose(): |
| 198 | """Test the propose() method of ExtractHiddenStatesProposer. |
| 199 | |
| 200 | This should: |
| 201 | 1. Accept target hidden states and sampled token IDs |
| 202 | 2. Return the sampled tokens as "draft" tokens (shape [batch_size, 1]) |
| 203 | 3. Cache the hidden states in the model's KV cache |
| 204 | """ |
| 205 | device = torch.device(DEVICE_TYPE) |
| 206 | |
| 207 | # Setup test parameters |
| 208 | batch_size = 2 |
| 209 | num_tokens = 5 |
| 210 | num_hidden_layers = 4 |
| 211 | |
| 212 | proposer = _create_proposer( |
| 213 | num_speculative_tokens=1, layer_ids=list(range(num_hidden_layers)) |
| 214 | ) |
| 215 | hidden_size = proposer.hidden_size |
| 216 | |
| 217 | # Create mock model |
| 218 | model_mock = mock.MagicMock() |
| 219 | proposer.model = model_mock |
| 220 | |
| 221 | # Mock attention layer names |
| 222 | proposer.attn_layer_names = ["cache_only_layers.28"] |
| 223 | |
| 224 | # Mock attention metadata builder |
| 225 | mock_attn_metadata = mock.MagicMock() |
| 226 | mock_attn_metadata_builder = mock.MagicMock() |
| 227 | mock_attn_metadata_builder.build_for_drafting.return_value = mock_attn_metadata |
| 228 | proposer.attn_metadata_builder = mock_attn_metadata_builder |
| 229 | |
| 230 | # Create input tensors |
| 231 | batch_spec = BatchSpec( |
| 232 | seq_lens=[3, 2], |
| 233 | query_lens=[3, 2], |
| 234 | ) |
| 235 | |
| 236 | common_attn_metadata = create_common_attn_metadata( |
| 237 | batch_spec, |
| 238 | block_size=16, |
| 239 | device=device, |
| 240 | ) |
| 241 | |
| 242 | # Create target hidden states: list of tensors, one per layer |
| 243 | # Each tensor has shape [num_tokens, hidden_size] |
| 244 | target_hidden_states = [ |
| 245 | torch.randn(num_tokens, hidden_size, dtype=proposer.dtype, device=device) |
| 246 | for _ in range(num_hidden_layers) |
| 247 | ] |
| 248 | |
| 249 | # Sampled token IDs from target model |
| 250 | sampled_token_ids = torch.tensor( |
| 251 | [42, 60], dtype=torch.int32, device=device |
| 252 | ).unsqueeze(-1) |
| 253 | |
| 254 | # Call propose |
| 255 | draft_tokens = proposer.propose( |
| 256 | num_speculative_tokens=1, |
| 257 | sampled_token_ids=sampled_token_ids, |
| 258 | target_hidden_states=target_hidden_states, |
| 259 | common_attn_metadata=common_attn_metadata, |
| 260 | slot_mappings=None, |
| 261 | ) |
| 262 | |
| 263 | # Verify draft tokens match sampled tokens |
| 264 | # Shape should be [batch_size, 1] for num_speculative_tokens=1 |
| 265 | assert draft_tokens.shape == (batch_size, 1) |
| 266 | assert torch.equal(draft_tokens, sampled_token_ids) |
| 267 | |
| 268 | # Verify the model was called |
| 269 | model_mock.assert_called_once() |
| 270 | |
| 271 | # Verify hidden states were copied to the buffer The stacked hidden states |
| 272 | # should have shape [num_tokens, num_hidden_layers, hidden_size] |
| 273 | expected_stacked = torch.stack(target_hidden_states, dim=1) |
| 274 | assert torch.allclose( |
| 275 | proposer.hidden_states[:num_tokens], expected_stacked, atol=1e-6 |
| 276 | ) |
| 277 | |
| 278 | |
| 279 | @pytest.mark.parametrize("num_hidden_layers", [1, 4, 8]) |
| 280 | def test_propose_different_layer_counts(num_hidden_layers): |
| 281 | """Test that propose works correctly with different numbers of hidden layers.""" |
| 282 | device = torch.device(DEVICE_TYPE) |
| 283 | |
| 284 | batch_size = 2 |
| 285 | num_tokens = 5 |
| 286 | |
| 287 | proposer = _create_proposer( |
| 288 | num_speculative_tokens=1, layer_ids=list(range(num_hidden_layers)) |
| 289 | ) |
| 290 | hidden_size = proposer.hidden_size |
| 291 | |
| 292 | # Setup mocks |
| 293 | model_mock = mock.MagicMock() |
| 294 | proposer.model = model_mock |
| 295 | proposer.attn_layer_names = ["cache_only_layers.28"] |
| 296 | |
| 297 | mock_attn_metadata_builder = mock.MagicMock() |
| 298 | mock_attn_metadata_builder.build_for_drafting.return_value = mock.MagicMock() |
| 299 | proposer.attn_metadata_builder = mock_attn_metadata_builder |
| 300 | |
| 301 | batch_spec = BatchSpec( |
| 302 | seq_lens=[3, 2], |
| 303 | query_lens=[3, 2], |
| 304 | ) |
| 305 | |
| 306 | common_attn_metadata = create_common_attn_metadata( |
| 307 | batch_spec, |
| 308 | block_size=16, |
| 309 | device=device, |
| 310 | ) |
| 311 | |
| 312 | # Create target hidden states |
| 313 | target_hidden_states = [ |
| 314 | torch.randn(num_tokens, hidden_size, dtype=proposer.dtype, device=device) |
| 315 | for _ in range(num_hidden_layers) |
| 316 | ] |
| 317 | |
| 318 | sampled_token_ids = torch.tensor( |
| 319 | [42, 60], dtype=torch.int32, device=device |
| 320 | ).unsqueeze(-1) |
| 321 | |
| 322 | draft_tokens = proposer.propose( |
| 323 | num_speculative_tokens=1, |
| 324 | sampled_token_ids=sampled_token_ids, |
| 325 | target_hidden_states=target_hidden_states, |
| 326 | common_attn_metadata=common_attn_metadata, |
| 327 | slot_mappings=None, |
| 328 | ) |
| 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) |