wehub-resource-sync
122 lines
3.7 KiB
Python
Executable File
122 lines
3.7 KiB
Python
Executable File
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Integration test for custom proposer class in speculative decoding.
|
|
|
|
Usage:
|
|
.venv/bin/python test_custom_proposer.py
|
|
"""
|
|
|
|
import os
|
|
|
|
import torch
|
|
|
|
from vllm import LLM, SamplingParams
|
|
from vllm.config import VllmConfig
|
|
|
|
MODEL_ID = "facebook/opt-125m"
|
|
NUM_SPEC_TOKENS = 5
|
|
|
|
|
|
class DummyDraftProposer:
|
|
"""Custom proposer class that repeats the last token of each sequence.
|
|
|
|
This demonstrates the class-based custom proposer interface.
|
|
"""
|
|
|
|
def __init__(self, vllm_config: VllmConfig):
|
|
"""Initialize the custom proposer.
|
|
|
|
Args:
|
|
vllm_config: vLLM configuration containing model and speculative settings.
|
|
"""
|
|
self.num_speculative_tokens = (
|
|
vllm_config.speculative_config.num_speculative_tokens
|
|
)
|
|
self.max_model_len = vllm_config.model_config.max_model_len
|
|
print(
|
|
f"[DummyDraftProposer.__init__] num_speculative_tokens="
|
|
f"{self.num_speculative_tokens}, max_model_len={self.max_model_len}"
|
|
)
|
|
|
|
def propose(
|
|
self,
|
|
sampled_token_ids: list[list[int]],
|
|
num_tokens_no_spec: int,
|
|
token_ids_cpu: torch.Tensor,
|
|
slot_mappings: torch.Tensor | None = None,
|
|
) -> list[list[int]]:
|
|
"""Generate draft tokens by repeating the last token of each sequence.
|
|
|
|
Args:
|
|
sampled_token_ids: Recently sampled token IDs per request.
|
|
num_tokens_no_spec: Number of non-speculative tokens per request.
|
|
token_ids_cpu: Full token IDs tensor on CPU.
|
|
slot_mappings: Slot mapping for KV cache (optional).
|
|
|
|
Returns:
|
|
List of draft token sequences for each request.
|
|
"""
|
|
# Cross-process flag to prove this method was executed
|
|
with open("proposer_called.flag", "w") as f:
|
|
f.write("called")
|
|
|
|
batch_size = len(sampled_token_ids)
|
|
last_tokens = [seq[-1] for seq in sampled_token_ids]
|
|
drafts = [[t] * self.num_speculative_tokens for t in last_tokens]
|
|
print(
|
|
f"[DummyDraftProposer.propose] batch_size={batch_size}, "
|
|
f"num_speculative_tokens={self.num_speculative_tokens}, "
|
|
f"drafts_shape={len(drafts)}x{len(drafts[0])}"
|
|
)
|
|
return drafts
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("=" * 60)
|
|
print("Custom Proposer Backend Integration Test")
|
|
print("=" * 60)
|
|
|
|
# Cleanup any leftover flag from previous failed runs
|
|
if os.path.exists("proposer_called.flag"):
|
|
os.remove("proposer_called.flag")
|
|
|
|
llm = LLM(
|
|
model=MODEL_ID,
|
|
speculative_config={
|
|
"model": f"{__name__}.DummyDraftProposer",
|
|
"num_speculative_tokens": NUM_SPEC_TOKENS,
|
|
},
|
|
gpu_memory_utilization=0.4,
|
|
enforce_eager=True,
|
|
)
|
|
|
|
prompts = [
|
|
"Hello, my name is",
|
|
"The future of AI is",
|
|
]
|
|
|
|
sampling_params = SamplingParams(
|
|
max_tokens=32,
|
|
temperature=0.0,
|
|
)
|
|
|
|
print(f"\nRunning generate with {len(prompts)} prompt(s)...\n")
|
|
outputs = llm.generate(prompts, sampling_params)
|
|
|
|
for output in outputs:
|
|
prompt = output.prompt
|
|
generated = output.outputs[0].text
|
|
print(f"Prompt: {prompt!r}")
|
|
print(f"Generated text: {generated!r}")
|
|
print("-" * 60)
|
|
|
|
# Verify the custom proposer's propose() was actually called across processes
|
|
assert os.path.exists("proposer_called.flag"), (
|
|
"The custom proposer's propose() method was never called!"
|
|
)
|
|
os.remove("proposer_called.flag")
|
|
|
|
print("✓ Custom proposer was actively used during generation!")
|
|
print("Test completed successfully.")
|