| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | |
| 6 | from vllm import TokensPrompt |
| 7 | from vllm.config import PoolerConfig |
| 8 | |
| 9 | |
| 10 | @pytest.mark.parametrize( |
| 11 | "model", |
| 12 | ["Qwen/Qwen3-0.6B"], |
| 13 | ) |
| 14 | @torch.inference_mode |
| 15 | def test_extract_hidden_states(hf_runner, vllm_runner, model: str): |
| 16 | n_prompt_tokens = [55, 56, 57] |
| 17 | token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens] |
| 18 | |
| 19 | with vllm_runner( |
| 20 | model, |
| 21 | max_model_len=128, |
| 22 | enforce_eager=True, |
| 23 | runner="pooling", |
| 24 | pooler_config=PoolerConfig(task="token_embed"), |
| 25 | enable_prefix_caching=True, |
| 26 | ) as vllm_model: |
| 27 | pooling_outputs = vllm_model.llm.encode( |
| 28 | [TokensPrompt(prompt_token_ids=t) for t in token_prompts], |
| 29 | pooling_task="token_embed", |
| 30 | ) |
| 31 | |
| 32 | for n, output in zip(n_prompt_tokens, pooling_outputs): |
| 33 | assert len(output.prompt_token_ids) == n |
| 34 | assert len(output.outputs.data) == n |
| 35 | assert output.num_cached_tokens == 0 |
| 36 | |
| 37 | # test enable_prefix_caching plus all pooling |
| 38 | # we need to skip reading cache at this request by |
| 39 | # request.skip_reading_prefix_cache |
| 40 | pooling_outputs = vllm_model.llm.encode( |
| 41 | [TokensPrompt(prompt_token_ids=t) for t in token_prompts], |
| 42 | pooling_task="token_embed", |
| 43 | ) |
| 44 | |
| 45 | for n, output in zip(n_prompt_tokens, pooling_outputs): |
| 46 | assert len(output.prompt_token_ids) == n |
| 47 | assert len(output.outputs.data) == n |
| 48 | assert output.num_cached_tokens == 0 |