[2/n] DP Enhancement: Reduce the memory usage of gathered buffer in DP attention

[ch-wan]

Stack from ghstack (oldest at bottom):

• [4/n] DP Enhancement: Optimize communication when dp < tp by using all_gather_into_tensor and reduce_scatter_tensor #8279
• [3/n] DP Enhancement: Padding tokens to max length when workload is balanced #8278
• -> [2/n] DP Enhancement: Reduce the memory usage of gathered buffer in DP attention #8277
• [1/n] DP Enhancement: Multiple minor refactor for eagle #8276

Two optimizations:

• batch.gathered_buffer is not needed when cuda graph is enabled. This PR defers the allocation of gathered_buffer after checking if cuda graph is runnable.
• Correcting global_num_tokens_for_logprob during draft extend to reduce the buffer size in logits processor.

This PR also pads input tokens to num_tokens * (speculative_num_steps + 1) during draft decode to avoid potential illegal memory access when copying hidden states to gathered buffer..

[gemini-code-assist]

Summary of Changes

Hello @ch-wan, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly optimizes memory usage within the distributed attention mechanism by implementing a more dynamic and precise allocation strategy for the gathered_buffer and related tensors. These changes are particularly impactful for speculative decoding scenarios. Additionally, the PR includes general code clean-up, type hinting improvements, and necessary adjustments to existing test cases to reflect the updated behavior and performance characteristics.

Highlights

• Memory Optimization for DP Attention: The primary goal of this PR is to reduce memory usage of the gathered_buffer in distributed attention (DP attention). This is achieved by removing redundant padding logic in _dp_gather and dp_scatter and by dynamically allocating the buffer.
• Dynamic Buffer Preparation: The gathered_buffer is now initialized and prepared dynamically via a new prepare_mlp_sync_batch method within ForwardBatch. This ensures memory is allocated more precisely when needed, especially during speculative decoding's draft_extend mode, and also handles padding of related tensors like input_ids and hidden_states.
• Refined Token Count Calculation: The calculation of global token counts (global_num_tokens, global_num_tokens_for_logprob) has been updated to leverage detailed speculative decoding information (from EagleDraftInput and EagleVerifyInput), leading to more accurate sizing and preventing over-allocation.
• Code Clean-up and Type Hinting: Several files received minor refactorings, including updating type hints from torch.tensor to torch.Tensor, moving imports into TYPE_CHECKING blocks for better performance and clarity, and removing a deprecated spec_num_draft_tokens field from ModelWorkerBatch.
• Test Adjustments: Existing hybrid DP/EP/TP/MTP tests were updated. This includes renaming numerous test classes, removing the test_mgsm_en methods, and adjusting the assertion thresholds for mmlu evaluation metrics to 0.48.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request effectively reduces memory usage by deferring the allocation of the gathered_buffer in DP attention until it's actually needed for MLP sync. The refactoring to use spec_info for speculative decoding parameters instead of spec_num_draft_tokens also improves code clarity and maintainability.

I've left a few comments, mainly concerning the significant changes in the test file test/srt/test_hybrid_dp_ep_tp_mtp.py, where a test case was removed and classes were renamed. I've also noted some good practices like adding and correcting type hints.

Overall, the changes look solid and well-aligned with the goal of memory optimization.