EAGLE Support DP>1

[Flechman]

Purpose

When using DP>1, the set_forward_context gathers the number of tokens across all DP ranks to get static batch size for all ranks for CudaGRAPH EP compatibility.
Even though CudaGRAPH is not yet implemented for the newly padded-speculative design, this DP-all-reduce/gather already serves for the target model and will serve for the draft model later.

When using EAGLE and DP>1 with num_speculative_tokens=k, and sending a single request, one DP rank will handle that request (going through propose()) and all other DP ranks will do a dummy forward. The dummy forward in the eagle module only performs a single forward-pass (so single DP all-reduce via set_forward_context), and we get k eagle forward-passes during propose(). This leads to a "phase shift" between the DP rank doing propose() (group A) and the DP ranks doing the dummy forward (group B): first step in group A will DP-all-reduce k times, and group B will DP-all-reduce a single time but doing k steps. If the finish-sync happens every k steps (in our case finish-sync happens every 32 steps), group B will enter a finish-sync all-reduce while group A is still doing the eagle DP-all-reduce because it's only entering step 2. This results in hanging.

Test Plan

Add an e2e correctness test with eagle and DP=2.

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[ekagra-ranjan]

Nice work! Can you also add a quality correctness test by using test_eagle_correctness() in DP > 1 setting? The added test just checks if the process hangs or not.

[Flechman]

Nice work! Can you also add a quality correctness test by using test_eagle_correctness() in DP > 1 setting? The added test just checks if the process hangs or not.

Thank you for pointing this out! I've just modified the test to compare outputs with an engine without eagle.

[Flechman]

Hey @benchislett @ekagra-ranjan @luccafong, any updates on this ?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Flechman.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[youkaichao]

cc @benchislett @luccafong

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Flechman.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Flechman]

@benchislett this one is ready for review 🙏

[benchislett]

LGTM, thanks!

[benchislett]

@Flechman please fix the broken test.