[Spec Decode] Disable Log Prob serialization to CPU for spec decoding for both draft and target models. (#6485)

2024-07-20 23:58:58 -07:00
parent d7f4178dd9
commit 14f91fe67c
8 changed files with 333 additions and 64 deletions
--- a/tests/spec_decode/test_dynamic_spec_decode.py
+++ b/tests/spec_decode/test_dynamic_spec_decode.py
@@ -32,6 +32,7 @@ def test_disable_spec_tokens(queue_size: int, batch_size: int, k: int,
                              scorer_worker=target_worker,
                              spec_decode_sampler=mock_spec_decode_sampler(
                                  acceptance_sampler_method),
+                              disable_logprobs=False,
                              metrics_collector=metrics_collector,
                              disable_by_batch_size=disable_by_batch_size)