[wangxiyuan]

expand_kernel ascend implementation by triton register, instead of patch way.

• vLLM version: v0.20.2
• vLLM main: vllm-project/vllm@0d4d334

[gemini-code-assist]

Summary of Changes

Hello, 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 enables platform-specific optimizations for the vLLM-Ascend project by integrating custom Triton kernels. By leveraging the upstream vLLM triton_dispatcher, the changes allow for more efficient execution of expansion operations on NPU hardware, ensuring better alignment with Ascend-specific performance requirements.

Highlights

• Triton Kernel Registration: Introduced a new registration module to integrate Ascend-specific Triton kernel implementations using the vLLM dispatcher mechanism.
• Expand Kernel Optimization: Added an Ascend-optimized version of the expand kernel that utilizes custom grid and block size calculations for improved NPU performance.
• Rejection Sampler Update: Disabled the legacy expand_batch_to_tokens assignment in the rejection sampler to favor the new dispatcher-based approach.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the 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 counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

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 introduces the registration of Ascend-specific Triton kernels using the triton_dispatcher mechanism, specifically for the expand operation. This change aims to replace manual function patching with a more integrated approach. Feedback indicates that the PR title and summary should be updated to follow the repository's style guide. Additionally, it is recommended to remove the commented-out code in patch_rejection_sampler.py instead of just disabling it, and to verify that the upstream implementation correctly utilizes the dispatcher for the 'expand' kernel to ensure the optimized implementation is actually used.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.