| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | import vllm |
| 8 | from vllm.lora.request import LoRARequest |
| 9 | from vllm.platforms import current_platform |
| 10 | |
| 11 | from ..utils import multi_gpu_test |
| 12 | |
| 13 | MODEL_PATH = "openai/gpt-oss-20b" |
| 14 | |
| 15 | PROMPT_TEMPLATE = """<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI. |
| 16 | Knowledge cutoff: 2024-06 |
| 17 | Current date: 2025-10-29 |
| 18 | |
| 19 | Reasoning: medium |
| 20 | |
| 21 | # Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>user<|message|>I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request. |
| 22 | " |
| 23 | ##Instruction: |
| 24 | farm contains tables such as city, farm, farm_competition, competition_record. Table city has columns such as City_ID, Official_Name, Status, Area_km_2, Population, Census_Ranking. City_ID is the primary key. |
| 25 | Table farm has columns such as Farm_ID, Year, Total_Horses, Working_Horses, Total_Cattle, Oxen, Bulls, Cows, Pigs, Sheep_and_Goats. Farm_ID is the primary key. |
| 26 | Table farm_competition has columns such as Competition_ID, Year, Theme, Host_city_ID, Hosts. Competition_ID is the primary key. |
| 27 | Table competition_record has columns such as Competition_ID, Farm_ID, Rank. Competition_ID is the primary key. |
| 28 | The Host_city_ID of farm_competition is the foreign key of City_ID of city. |
| 29 | The Farm_ID of competition_record is the foreign key of Farm_ID of farm. |
| 30 | The Competition_ID of competition_record is the foreign key of Competition_ID of farm_competition. |
| 31 | |
| 32 | |
| 33 | ###Input: |
| 34 | {context} |
| 35 | |
| 36 | ###Response:<|end|><|start|>assistant<|channel|>final<|message|>""" # noqa: E501 |
| 37 | |
| 38 | EXPECTED_LORA_OUTPUT = [ |
| 39 | "SELECT avg(Working_Horses) FROM farm WHERE Total_Horses > 5000", |
| 40 | "SELECT max(Cows) , min(Cows) FROM farm", |
| 41 | "SELECT max(Cows) , min(Cows) FROM farm", |
| 42 | ] |
| 43 | |
| 44 | |
| 45 | def reformat(text: str) -> str: |
| 46 | # Remove all spaces immediately before or after comma |
| 47 | text = ",".join(map(str.strip, text.split(","))) |
| 48 | # Remove duplicated blank spaces |
| 49 | text = " ".join(map(str.strip, text.split())) |
| 50 | return text |
| 51 | |
| 52 | |
| 53 | def _enable_deterministic_lora_shrink(monkeypatch: pytest.MonkeyPatch) -> None: |
| 54 | if not current_platform.is_rocm(): |
| 55 | return |
| 56 | |
| 57 | # These tests assert exact greedy outputs. Force the Triton LoRA shrink |
| 58 | # kernel to use SPLIT_K=1 so it stores the complete reduction directly |
| 59 | # instead of accumulating split-K partial results with atomic_add. |
| 60 | monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") |
| 61 | # The kernel configuration reads VLLM_BATCH_INVARIANT at import time. |
| 62 | # Spawn the engine process so it observes this setting even if the LoRA |
| 63 | # Triton utilities were already imported during test collection. |
| 64 | monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") |
| 65 | |
| 66 | |
| 67 | def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None: |
| 68 | prompts = [ |
| 69 | PROMPT_TEMPLATE.format( |
| 70 | context="Give the average number of working horses on farms with more than 5000 total horses." # noqa: E501 |
| 71 | ), # noqa: E501 |
| 72 | PROMPT_TEMPLATE.format( |
| 73 | context="What are the maximum and minimum number of cows across all farms." |
| 74 | ), |
| 75 | PROMPT_TEMPLATE.format( |
| 76 | context="Return the maximum and minimum number of cows across all farms." |
| 77 | ), |
| 78 | ] |
| 79 | sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64) |
| 80 | outputs = llm.generate( |
| 81 | prompts, |
| 82 | sampling_params, |
| 83 | lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None, |
| 84 | ) |
| 85 | # Print the outputs. |
| 86 | generated_texts: list[str] = [] |
| 87 | for output in outputs: |
| 88 | prompt = output.prompt |
| 89 | generated_text = output.outputs[0].text.strip() |
| 90 | generated_texts.append(generated_text) |
| 91 | print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") |
| 92 | for i in range(len(EXPECTED_LORA_OUTPUT)): |
| 93 | # The generated text may have different numbers of blank space, |
| 94 | # so reformat to compare. |
| 95 | compactGeneratedStr = reformat(generated_texts[i]) |
| 96 | compactExpectedStr = reformat(EXPECTED_LORA_OUTPUT[i]) |
| 97 | if not generated_texts[i].startswith( |
| 98 | EXPECTED_LORA_OUTPUT[i] |
| 99 | ) and not compactGeneratedStr.startswith(compactExpectedStr): |
| 100 | raise AssertionError( |
| 101 | f"Generated: {generated_texts[i]}, Expected: {EXPECTED_LORA_OUTPUT[i]}" |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | @pytest.mark.parametrize( |
| 106 | "mxfp4_use_marlin", |
| 107 | [ |
| 108 | False, |
| 109 | pytest.param( |
| 110 | True, |
| 111 | marks=pytest.mark.skipif( |
| 112 | current_platform.is_rocm() or current_platform.is_xpu(), |
| 113 | reason="marlin not supported", |
| 114 | ), |
| 115 | ), |
| 116 | ], |
| 117 | ) |
| 118 | @pytest.mark.parametrize("specialize_active_lora", [True, False]) |
| 119 | def test_gpt_oss_lora( |
| 120 | gptoss20b_lora_files, |
| 121 | mxfp4_use_marlin, |
| 122 | specialize_active_lora, |
| 123 | monkeypatch: pytest.MonkeyPatch, |
| 124 | ): |
| 125 | _enable_deterministic_lora_shrink(monkeypatch) |
| 126 | |
| 127 | llm = vllm.LLM( |
| 128 | MODEL_PATH, |
| 129 | max_model_len=1024, |
| 130 | enable_lora=True, |
| 131 | max_loras=4, |
| 132 | max_lora_rank=8, |
| 133 | max_num_seqs=2, |
| 134 | max_num_batched_tokens=2048, |
| 135 | specialize_active_lora=specialize_active_lora, |
| 136 | moe_backend="marlin" if mxfp4_use_marlin else "auto", |
| 137 | linear_backend="marlin" if mxfp4_use_marlin else "auto", |
| 138 | compilation_config=vllm.config.CompilationConfig( # Avoid OOM |
| 139 | cudagraph_specialize_lora=False, |
| 140 | ), |
| 141 | ) |
| 142 | |
| 143 | generate_and_test(llm, gptoss20b_lora_files, lora_id=1) |
| 144 | generate_and_test(llm, gptoss20b_lora_files, lora_id=2) |
| 145 | |
| 146 | |
| 147 | @multi_gpu_test(num_gpus=2) |
| 148 | @pytest.mark.parametrize("fully_sharded_loras", [False, True]) |
| 149 | @pytest.mark.parametrize( |
| 150 | "mxfp4_use_marlin", |
| 151 | [ |
| 152 | False, |
| 153 | pytest.param( |
| 154 | True, |
| 155 | marks=pytest.mark.skipif( |
| 156 | current_platform.is_rocm() or current_platform.is_xpu(), |
| 157 | reason="marlin not supported", |
| 158 | ), |
| 159 | ), |
| 160 | ], |
| 161 | ) |
| 162 | def test_gpt_oss_lora_tp2( |
| 163 | gptoss20b_lora_files, |
| 164 | fully_sharded_loras, |
| 165 | mxfp4_use_marlin, |
| 166 | monkeypatch: pytest.MonkeyPatch, |
| 167 | ): |
| 168 | _enable_deterministic_lora_shrink(monkeypatch) |
| 169 | |
| 170 | llm = vllm.LLM( |
| 171 | MODEL_PATH, |
| 172 | max_model_len=1024, |
| 173 | enable_lora=True, |
| 174 | max_loras=2, |
| 175 | max_num_seqs=2, |
| 176 | max_num_batched_tokens=2048, |
| 177 | tensor_parallel_size=2, |
| 178 | gpu_memory_utilization=0.8, |
| 179 | fully_sharded_loras=fully_sharded_loras, |
| 180 | enable_expert_parallel=not fully_sharded_loras, |
| 181 | moe_backend="marlin" if mxfp4_use_marlin else "auto", |
| 182 | linear_backend="marlin" if mxfp4_use_marlin else "auto", |
| 183 | compilation_config=vllm.config.CompilationConfig( |
| 184 | cudagraph_specialize_lora=False, |
| 185 | ), |
| 186 | ) |
| 187 | |
| 188 | generate_and_test(llm, gptoss20b_lora_files, lora_id=1) |
| 189 | generate_and_test(llm, gptoss20b_lora_files, lora_id=2) |