[Enhancement] Custom Logit Processor Improvement (#2998)

Signed-off-by: Hongpeng Guo <hpguo@anyscale.com>
2025-01-20 02:00:35 -08:00
parent 2584f6d944
commit 583697cd71
6 changed files with 79 additions and 28 deletions
@@ -966,6 +966,7 @@ class Scheduler:
            self.model_config,
            self.enable_overlap,
            self.spec_algorithm,
+            self.server_args.enable_custom_logit_processor,
        )
        new_batch.prepare_for_extend()

@@ -1520,6 +1521,7 @@ class Scheduler:
            self.model_config,
            self.enable_overlap,
            self.spec_algorithm,
+            self.server_args.enable_custom_logit_processor,
        )
        idle_batch.prepare_for_idle()
        return idle_batch