A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 00:30 UTC
[CI] Fix test_lora_with_spec_decode on V2 model runner (#43314)
Signed-off-by: haosdent <haosdent@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
haosdent committed
May 22, 2026 at 14:24 UTC
5ea76fa89aa2e307f0d9a2e7fc19d13aed65a82f
2 files changed
+53
-36
tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py
+40
-35
@@ -67,6 +67,12 @@ def test_batch_inference_correctness(
67
68
method, model_name, spec_model_name, lora_path, tp_size = model_setup
69
70
+ prompts = [LORA_TEST_PROMPT_MAP[lora_path]] * 100
71
+ lora_request = LoRARequest("adapter", 1, lora_path)
72
+ sampling_params = SamplingParams(
73
+ temperature=0.0, top_p=1.0, top_k=-1, seed=SEED, max_tokens=128
74
+ )
75
+
76
# without speculative decoding
77
ref_llm = LLM(
78
model=model_name,
@@ -79,19 +85,14 @@ def test_batch_inference_correctness(
85
max_cpu_loras=1,
86
max_lora_rank=16,
87
)
82
-
83
- prompts = [LORA_TEST_PROMPT_MAP[lora_path]] * 100
84
- lora_request = LoRARequest("adapter", 1, lora_path)
85
- sampling_params = SamplingParams(
86
- temperature=0.0, top_p=1.0, top_k=-1, seed=SEED, max_tokens=128
87
- )
88
-
89
- ref_outputs = ref_llm.generate(
90
- prompts, sampling_params, lora_request=lora_request
91
- )
92
- del ref_llm
93
- torch.accelerator.empty_cache()
94
- cleanup_dist_env_and_memory()
88
+ try:
89
+ ref_outputs = ref_llm.generate(
90
+ prompts, sampling_params, lora_request=lora_request
91
+ )
92
+ finally:
93
+ del ref_llm
94
+ torch.accelerator.empty_cache()
95
+ cleanup_dist_env_and_memory()
96
97
lora_spec_llm = LLM(
98
model=model_name,
@@ -110,25 +111,29 @@ def test_batch_inference_correctness(
111
max_cpu_loras=1,
112
max_lora_rank=16,
113
)
113
-
114
- lora_spec_outputs = lora_spec_llm.generate(
115
- prompts, sampling_params, lora_request=lora_request
116
- )
117
-
118
- matches = 0
119
- misses = 0
120
- for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
121
- if ref_output.outputs[0].text == spec_output.outputs[0].text:
122
- matches += 1
123
- else:
124
- misses += 1
125
- print(f"ref_output: {ref_output.outputs[0].text}")
126
- print(f"spec_output: {spec_output.outputs[0].text}")
127
-
128
- # Heuristic: expect at least 90% of the prompts to match exactly
129
- # Upon failure, inspect the outputs to check for inaccuracy.
130
- print(f"match ratio: {matches}/{len(ref_outputs)}")
131
- assert matches > int(0.90 * len(ref_outputs))
132
- del lora_spec_llm
133
- torch.accelerator.empty_cache()
134
- cleanup_dist_env_and_memory()
114
+ try:
115
+ lora_spec_outputs = lora_spec_llm.generate(
116
+ prompts, sampling_params, lora_request=lora_request
117
+ )
118
+
119
+ matches = 0
120
+ for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
121
+ if ref_output.outputs[0].text == spec_output.outputs[0].text:
122
+ matches += 1
123
+ else:
124
+ print(f"ref_output: {ref_output.outputs[0].text}")
125
+ print(f"spec_output: {spec_output.outputs[0].text}")
126
+
127
+ # Heuristic threshold: under greedy verification, the spec-decode
128
+ # output should equal the non-spec output (modulo FP noise from the
129
+ # target's verify-path matmul running at seqlen
130
+ # num_speculative_tokens+1 vs 1). 90% leaves slack for that noise.
131
+ threshold = int(0.90 * len(ref_outputs))
132
+ print(f"match ratio: {matches}/{len(ref_outputs)}")
133
+ assert matches > threshold, (
134
+ f"match ratio {matches}/{len(ref_outputs)} <= {threshold}"
135
+ )
136
+ finally:
137
+ del lora_spec_llm
138
+ torch.accelerator.empty_cache()
139
+ cleanup_dist_env_and_memory()
vllm/v1/worker/gpu/model_runner.py
+13
-1
@@ -1091,13 +1091,25 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1091
# FIXME: Replace this with LoRA warmup:
1092
# https://github.com/vllm-project/vllm/pull/35536
1093
assert hasattr(self, "lora_manager")
1094
- self.lora_manager._adapter_manager.set_adapter_mapping(
1094
+ adapter_manager = self.lora_manager._adapter_manager
1095
+ adapter_manager.set_adapter_mapping(
1096
LoRAMapping(
1097
index_mapping=(0,) * input_batch.num_tokens_after_padding,
1098
prompt_mapping=(0,) * input_batch.num_reqs,
1099
is_prefill=True,
1100
)
1101
)
1102
+ seen_wrappers: set[int] = set()
1103
+ for punica_wrapper in adapter_manager.punica_wrapper_mapping.values():
1104
+ if id(punica_wrapper) in seen_wrappers:
1105
+ continue
1106
+ seen_wrappers.add(id(punica_wrapper))
1107
+ for kernel_meta in (
1108
+ punica_wrapper.token_mapping_meta, # type: ignore[attr-defined]
1109
+ punica_wrapper.prompt_mapping_meta, # type: ignore[attr-defined]
1110
+ ):
1111
+ kernel_meta.no_lora_flag_cpu[0] = False
1112
+ kernel_meta.num_active_loras_cpu[0] = 1
1113
1114
attn_metadata = None
1115
slot_mappings_by_layer = None