| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Integration test for custom proposer class in speculative decoding. |
| 5 | |
| 6 | Usage: |
| 7 | .venv/bin/python test_custom_proposer.py |
| 8 | """ |
| 9 | |
| 10 | import os |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | from vllm import LLM, SamplingParams |
| 15 | from vllm.config import VllmConfig |
| 16 | |
| 17 | MODEL_ID = "facebook/opt-125m" |
| 18 | NUM_SPEC_TOKENS = 5 |
| 19 | |
| 20 | |
| 21 | class DummyDraftProposer: |
| 22 | """Custom proposer class that repeats the last token of each sequence. |
| 23 | |
| 24 | This demonstrates the class-based custom proposer interface. |
| 25 | """ |
| 26 | |
| 27 | def __init__(self, vllm_config: VllmConfig): |
| 28 | """Initialize the custom proposer. |
| 29 | |
| 30 | Args: |
| 31 | vllm_config: vLLM configuration containing model and speculative settings. |
| 32 | |
| 33 | """ |
| 34 | assert vllm_config.speculative_config is not None |
| 35 | self.num_speculative_tokens = ( |
| 36 | vllm_config.speculative_config.num_speculative_tokens |
| 37 | ) |
| 38 | self.max_model_len = vllm_config.model_config.max_model_len |
| 39 | print( |
| 40 | f"[DummyDraftProposer.__init__] num_speculative_tokens=" |
| 41 | f"{self.num_speculative_tokens}, max_model_len={self.max_model_len}" |
| 42 | ) |
| 43 | |
| 44 | def propose( |
| 45 | self, |
| 46 | sampled_token_ids: list[list[int]], |
| 47 | num_tokens_no_spec: int, |
| 48 | token_ids_cpu: torch.Tensor, |
| 49 | slot_mappings: torch.Tensor | None = None, |
| 50 | ) -> list[list[int]]: |
| 51 | """Generate draft tokens by repeating the last token of each sequence. |
| 52 | |
| 53 | Args: |
| 54 | sampled_token_ids: Recently sampled token IDs per request. |
| 55 | num_tokens_no_spec: Number of non-speculative tokens per request. |
| 56 | token_ids_cpu: Full token IDs tensor on CPU. |
| 57 | slot_mappings: Slot mapping for KV cache (optional). |
| 58 | |
| 59 | Returns: |
| 60 | List of draft token sequences for each request. |
| 61 | |
| 62 | """ |
| 63 | # Cross-process flag to prove this method was executed |
| 64 | with open("proposer_called.flag", "w") as f: |
| 65 | f.write("called") |
| 66 | |
| 67 | batch_size = len(sampled_token_ids) |
| 68 | last_tokens = [seq[-1] for seq in sampled_token_ids] |
| 69 | drafts = [[t] * self.num_speculative_tokens for t in last_tokens] |
| 70 | print( |
| 71 | f"[DummyDraftProposer.propose] batch_size={batch_size}, " |
| 72 | f"num_speculative_tokens={self.num_speculative_tokens}, " |
| 73 | f"drafts_shape={len(drafts)}x{len(drafts[0])}" |
| 74 | ) |
| 75 | return drafts |
| 76 | |
| 77 | |
| 78 | if __name__ == "__main__": |
| 79 | print("=" * 60) |
| 80 | print("Custom Proposer Backend Integration Test") |
| 81 | print("=" * 60) |
| 82 | |
| 83 | # Cleanup any leftover flag from previous failed runs |
| 84 | if os.path.exists("proposer_called.flag"): |
| 85 | os.remove("proposer_called.flag") |
| 86 | |
| 87 | llm = LLM( |
| 88 | model=MODEL_ID, |
| 89 | speculative_config={ |
| 90 | "model": f"{__name__}.DummyDraftProposer", |
| 91 | "num_speculative_tokens": NUM_SPEC_TOKENS, |
| 92 | }, |
| 93 | gpu_memory_utilization=0.4, |
| 94 | enforce_eager=True, |
| 95 | ) |
| 96 | |
| 97 | prompts = [ |
| 98 | "Hello, my name is", |
| 99 | "The future of AI is", |
| 100 | ] |
| 101 | |
| 102 | sampling_params = SamplingParams( |
| 103 | max_tokens=32, |
| 104 | temperature=0.0, |
| 105 | ) |
| 106 | |
| 107 | print(f"\nRunning generate with {len(prompts)} prompt(s)...\n") |
| 108 | outputs = llm.generate(prompts, sampling_params) |
| 109 | |
| 110 | for output in outputs: |
| 111 | prompt = output.prompt |
| 112 | generated = output.outputs[0].text |
| 113 | print(f"Prompt: {prompt!r}") |
| 114 | print(f"Generated text: {generated!r}") |
| 115 | print("-" * 60) |
| 116 | |
| 117 | # Verify the custom proposer's propose() was actually called across processes |
| 118 | assert os.path.exists("proposer_called.flag"), ( |
| 119 | "The custom proposer's propose() method was never called!" |
| 120 | ) |
| 121 | os.remove("proposer_called.flag") |
| 122 | |
| 123 | print("✓ Custom proposer was actively used during generation!") |
| 124 | print("Test completed successfully.") |