The condition self.spec_algo.is_eagle() or self.spec_algo.is_standalone() is repeated in this file and also in scheduler_output_processor_mixin.py. To improve maintainability, consider adding a helper method to the SpeculativeAlgorithm class in spec_info.py. For example:

def is_eagle_or_standalone(self):
    return self.is_eagle() or self.is_standalone()

This would centralize the logic and make the code cleaner.

The condition self.is_v2_eagle or self.is_v2_standalone is repeated in several places (here, maybe_wait_verify_done, and in scheduler.py and scheduler_output_processor_mixin.py). To improve maintainability and reduce code duplication, consider adding a new property to the ScheduleBatch class that encapsulates this logic. For example:

@property
def is_v2_spec(self):
    return self.is_v2_eagle or self.is_v2_standalone

Then you can simplify this condition to if self.is_v2_spec:. This would make the code cleaner and easier to modify in the future if more v2 speculative algorithms are added.

The hardcoded pp_rank=0 with a FIXME comment indicates that this worker may not function correctly with pipeline parallelism. This should be resolved to properly support pipeline parallel setups.

The __init__ method of StandaloneDraftWorker is an exact copy of EagleDraftWorker.__init__. This code duplication can be avoided by removing the __init__ method from StandaloneDraftWorker and letting it inherit from EagleDraftWorker. Since EagleDraftWorker.__init__ calls self.init_lm_head(), your override of init_lm_head will be correctly used, achieving the same goal with much cleaner code.

With the suggested is_v2_algo method in SpeculativeAlgorithm, this condition can be simplified.

With the suggested is_v2_algo method in SpeculativeAlgorithm, this condition can be simplified.

With the suggested is_v2_algo method in SpeculativeAlgorithm, this condition can be simplified.

Following the suggestion to add is_v2_algo to SpeculativeAlgorithm, you can introduce a corresponding is_v2_spec property here. This will centralize the logic for checking v2 speculative algorithms and simplify conditions in this and other files.

With the new is_v2_spec property on ScheduleBatch, this can be simplified.

With the new is_v2_spec property on ScheduleBatch, this can be simplified.

With the suggested is_v2_algo method in SpeculativeAlgorithm, this condition can be simplified.

With the new is_v2_spec property on ScheduleBatch, this can be simplified.

With the new is_v2_spec property on ScheduleBatch, this can be simplified.

To improve maintainability and reduce code duplication, you can introduce a new property is_v2_speculative_overlap that combines the checks for is_v2_eagle and is_v2_standalone. This new property can then be used here and in other places like maybe_wait_verify_done in this file, as well as in scheduler_output_processor_mixin.py, to simplify the conditional logic.

For better readability and conciseness, you can use the in operator with a set for this check.

                self.speculative_algorithm in {"EAGLE", "STANDALONE"}

The __init__ method of StandaloneDraftWorker is almost a complete copy of EagleDraftWorker.__init__. To avoid this large code duplication and improve maintainability, you can call super().__init__ and rely on the parent class's implementation. The only functional difference, which is not sharing the LM head, is already correctly handled by overriding init_lm_head.

The __init__ method of StandaloneWorkerV2 is a near-exact copy of EAGLEWorkerV2.__init__, with the only difference being the instantiation of StandaloneDraftWorker instead of EagleDraftWorker. To reduce code duplication and make the code more maintainable, you can call super().__init__ and then just override the self._draft_worker attribute.

    def __init__(
        self,
        server_args: ServerArgs,
        gpu_id: int,
        tp_rank: int,
        dp_rank: Optional[int],
        moe_ep_rank: int,
        nccl_port: int,
        target_worker: TpModelWorker,
    ):
        super().__init__(
            server_args,
            gpu_id,
            tp_rank,
            dp_rank,
            moe_ep_rank,
            nccl_port,
            target_worker,
        )
        # Create our custom draft worker that doesn't share embeddings/lm_head
        self._draft_worker = StandaloneDraftWorker(
            server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
        )