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 4, 08:45 UTC
[ROCm][CI] Align spec decode logprob test prefill settings (#41335)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
May 3, 2026 at 23:33 UTC
01d4d1ad375dc5854779c593eee093bcebb0cada
1 file changed
+16
-19
tests/v1/sample/test_logprobs.py
+16
-19
@@ -33,11 +33,10 @@ PROMPT = BatchLogprobsComposition.PROMPT
33
SAMPLE_PROMPT = BatchLogprobsComposition.SAMPLE_PROMPT
34
35
# On ROCm, floating-point reductions in attention and GEMM kernels are
36
-# non-associative and sensitive to batch geometry. The ref LLM (no spec
37
-# decode, default scheduling) and the spec-decode LLM (chunked prefill,
38
-# different effective batch sizes) follow different reduction orders,
39
-# producing numerically divergent logprobs that get misattributed to
40
-# spec-decode incorrectness.
36
+# non-associative and sensitive to batch geometry. If the ref LLM and
37
+# spec-decode LLM use different scheduling or batch geometry, they can
38
+# follow different reduction orders and produce numerically divergent
39
+# logprobs that get misattributed to spec-decode incorrectness.
40
#
41
# Force LLM instances into an identical, deterministic execution
42
# mode so the test isolates spec-decode correctness only:
@@ -1086,18 +1085,25 @@ def test_spec_decode_logprobs(
1085
)
1086
1087
max_model_len = 256
1089
-
1090
- # Run base LLM.
1091
- ref_llm = LLM(
1092
- model=model_name,
1088
+ llm_kwargs = dict(
1089
max_logprobs=5,
1090
max_model_len=max_model_len,
1091
seed=42,
1092
logprobs_mode=logprobs_mode,
1093
gpu_memory_utilization=0.4,
1094
+ # Force the same prefill chunking for both the base model and
1095
+ # spec decode model so the comparison isolates spec decode.
1096
+ enable_chunked_prefill=True,
1097
+ max_num_batched_tokens=32,
1098
enable_prefix_caching=False,
1099
**ROCM_DETERMINISM_KWARGS,
1100
)
1101
+
1102
+ # Run base LLM.
1103
+ ref_llm = LLM(
1104
+ model=model_name,
1105
+ **llm_kwargs,
1106
+ )
1107
ref_results = ref_llm.generate(
1108
[prompt, prompt], [sampling_params, penalty_sampling_params]
1109
)
@@ -1117,16 +1123,7 @@ def test_spec_decode_logprobs(
1123
spec_llm = LLM(
1124
model_name,
1125
speculative_config=spec_config_with_len,
1120
- max_logprobs=5,
1121
- max_model_len=max_model_len,
1122
- seed=42,
1123
- logprobs_mode=logprobs_mode,
1124
- gpu_memory_utilization=0.4,
1125
- # Force prefill chunking
1126
- enable_chunked_prefill=True,
1127
- max_num_batched_tokens=32,
1128
- enable_prefix_caching=False,
1129
- **ROCM_DETERMINISM_KWARGS,
1126
+ **llm_kwargs,
1127
)
1128
spec_results = spec_llm.generate(
1129
[prompt, prompt], [sampling_params, penalty_sampling_params]