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, 22:45 UTC
[CPU] Fix spec decode kernel signatures for synthetic mode compatibility (#41932)
Signed-off-by: jmamou <jonathan.mamou@intel.com> Signed-off-by: Jonathan Mamou <jonathan.mamou@intel.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>
Jonathan Mamou committed
May 10, 2026 at 15:07 UTC
a54f0d1049c3acf18047f21e157d6452b7305957
1 file changed
+15
vllm/utils/cpu_triton_utils.py
+15
@@ -204,8 +204,16 @@ def _rejection_greedy_sample_kernel_impl(
204
bonus_token_ids,
205
is_greedy,
206
max_spec_len,
207
+ uniform_probs=None,
208
+ synthetic_conditional_rates=None,
209
+ SYNTHETIC_MODE=False,
210
):
211
# C++ kernel expects int64 for all integer tensors.
212
+ # Note: uniform_probs, synthetic_conditional_rates, and SYNTHETIC_MODE are
213
+ # passed by the rejection sampler for synthetic mode support, but are not
214
+ # yet implemented in the C++ CPU kernel. We accept them here to maintain
215
+ # compatibility with the kernel calling convention.
216
+ assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling"
217
orig_dtype = output_token_ids.dtype
218
output_token_ids_i64 = _ensure_int64(output_token_ids)
219
torch.ops._C.rejection_greedy_sample_kernel_impl(
@@ -233,11 +241,18 @@ def _rejection_random_sample_kernel_impl(
241
is_greedy,
242
max_spec_len,
243
vocab_size,
244
+ synthetic_conditional_rates=None,
245
NO_DRAFT_PROBS=False,
246
+ SYNTHETIC_MODE=False,
247
):
248
# C++ kernel expects int64 for all integer tensors and float32 for probs.
249
# uniform_probs is intentionally float64 in Python to avoid exact-zero
250
# samples; cast to float32 here for C++ compatibility.
251
+ # Note: synthetic_conditional_rates and SYNTHETIC_MODE are passed by the
252
+ # rejection sampler for synthetic mode support, but are not yet implemented
253
+ # in the C++ CPU kernel. We accept them here to maintain compatibility with
254
+ # the kernel calling convention.
255
+ assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling"
256
orig_dtype = output_token_ids.dtype
257
output_token_ids_i64 = _ensure_int64(output_token_ids)
258
torch.ops._C.rejection_random_sample_kernel_impl(