wehub-resource-sync
85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import os
|
|
import tempfile
|
|
|
|
from vllm import LLM, SamplingParams
|
|
from vllm.config.kv_transfer import KVTransferConfig
|
|
from vllm.distributed.kv_transfer.kv_connector.v1 import (
|
|
example_hidden_states_connector,
|
|
)
|
|
|
|
# NOTE: If changing the interface of the ExampleHiddenStatesConnector, please also
|
|
# update the benchmark in benchmarks/benchmark_hidden_state_extraction.py
|
|
# and the docs in docs/features/speculative_decoding/extract_hidden_states.md
|
|
|
|
# Example: Using the custom "extract_hidden_states" speculator method and
|
|
# ExampleHiddenStatesConnector to extract and save hidden states from vllm
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdirname:
|
|
llm = LLM(
|
|
model="Qwen/Qwen3-8B", # Your target model
|
|
speculative_config={
|
|
"method": "extract_hidden_states",
|
|
"num_speculative_tokens": 1,
|
|
"draft_model_config": {
|
|
"hf_config": {
|
|
"eagle_aux_hidden_state_layer_ids": [ # Target model layer indices
|
|
1,
|
|
2,
|
|
3,
|
|
4,
|
|
],
|
|
},
|
|
},
|
|
},
|
|
kv_transfer_config=KVTransferConfig(
|
|
kv_connector="ExampleHiddenStatesConnector",
|
|
kv_role="kv_producer",
|
|
kv_connector_extra_config={
|
|
"shared_storage_path": tmpdirname,
|
|
"allow_custom_save_path": True,
|
|
},
|
|
),
|
|
)
|
|
|
|
prompts = ["Generate a sentence with hidden states", "Write a python function"]
|
|
|
|
# One request uses defaults, the other uses a custom save path and
|
|
# includes output token hidden states via per-request kv_transfer_params.
|
|
sampling_params_list = [
|
|
SamplingParams(max_tokens=1),
|
|
SamplingParams(
|
|
max_tokens=10,
|
|
extra_args={
|
|
"kv_transfer_params": {
|
|
"hidden_states_path": os.path.join(
|
|
tmpdirname, "custom_output.safetensors"
|
|
),
|
|
"include_output_tokens": True,
|
|
}
|
|
},
|
|
),
|
|
]
|
|
outputs = llm.generate(prompts, sampling_params_list)
|
|
|
|
for output in outputs:
|
|
print("\nPrompt:", output.prompt)
|
|
print("Prompt token ids:", output.prompt_token_ids)
|
|
|
|
hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
|
|
assert hidden_states_path is not None
|
|
print("Hidden states path:", hidden_states_path)
|
|
|
|
obj = example_hidden_states_connector.load_hidden_states(hidden_states_path)
|
|
token_ids = obj["token_ids"]
|
|
hidden_states = obj["hidden_states"]
|
|
|
|
print("Extracted token ids:", token_ids)
|
|
print(
|
|
"Extracted hidden states shape:", hidden_states.shape
|
|
) # [num_tokens, num_extracted_layers, hidden_size]
|
|
print("Extracted hidden states:", hidden_states)
|
|
|
|
example_hidden_states_connector.cleanup_hidden_states(hidden_states_path)
|