Skip to content

Commit 7b54002

Browse files
committed
Address DSpark fused sampler review nits
Document that the Python seed counter must not be frozen by any future CUDA graph capture around sampling. Avoid an unnecessary clone on the top-k/top-p fused sampler path when dtype conversion already produced a float32 copy, while preserving a clone for existing float32 inputs.
1 parent f4e0e08 commit 7b54002

1 file changed

Lines changed: 5 additions & 1 deletion

File tree

vllm/v1/spec_decode/dspark_sampling.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
# the same step. Determinism of the *draft* noise across runs is harmless: output
3131
# randomness still comes from the target's synchronized rejection sampler, and
3232
# the counter keeps advancing across requests so proposals are not repeated.
33+
# If sampling is ever CUDA-graph-captured, keep this seed outside replay or pass
34+
# it as graph input; otherwise the replayed graph would reuse one draft seed.
3335
_FUSED_MARKOV_SEED_COUNTER = 0x9E3779B1
3436

3537

@@ -335,7 +337,9 @@ def sample_dspark_markov_block_fused(
335337
sample_logits = step_logits
336338
sample_inv_temp = inv_temp
337339
if uses_top_k_top_p:
338-
sample_logits = step_logits.to(torch.float32).clone()
340+
sample_logits = step_logits.to(torch.float32)
341+
if sample_logits.dtype == step_logits.dtype:
342+
sample_logits = sample_logits.clone()
339343
sample_logits.mul_(inv_temp.view(-1, 1))
340344
sample_logits = apply_top_k_top_p(sample_logits, top_k, top_p)
341345
assert unit_inv_temp is not None

0 commit comments

Comments
 (0)