[Triton] Centralize Ascend extension op dispatch in triton_utils

[linfeng-yuan]

What this PR does / why we need it?

This pull request refactors the dispatch mechanism for the triton-ascend-specific operators insert_slice, extract_slice, and get_element to ensure compatibility with both CANN 8.5 and 9.0.

A unified helper function, _resolve_triton_ascend_op, has been introduced in vllm_ascend/ops/triton/triton_utils.py. This function dynamically resolves these operators by first attempting to import them from the triton.language.extra.cann.extension module, which is present in newer CANN versions. If that fails, it falls back to the standard triton.language module.

This approach centralizes operator dispatch logic, allowing individual Triton kernels to use these functions without being aware of the underlying Triton/CANN version. All call sites have been updated to use these new unified functions.

Does this PR introduce any user-facing change?

No. This is an internal refactoring of operator implementations and does not introduce any user-facing changes.

How was this patch tested?

CI is expected to pass with existing tests.

Testing Context:

• vLLM version: v0.16.0
• vLLM main: 15d76f74e2fdb12a95ea00f0ca283acf6219a2b7

[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.

[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 introduces a significant refactoring to how specific Triton operations, namely insert_slice, extract_slice, and get_element, are handled within the codebase. By creating a centralized resolution mechanism in triton_utils, the change enhances the adaptability and compatibility of Triton kernels with different versions of CANN, specifically 8.5 and 9.0. This ensures a more robust and future-proof integration of Triton-based operations.

Highlights

• Unified Triton Operations: The pull request centralizes the dispatch of insert_slice, extract_slice, and get_element Triton operations through a new utility in triton_utils.py.
• CANN Compatibility: This refactoring aims to improve compatibility with CANN 8.5 and 9.0 by first attempting to load these operations from triton.language.extra.cann.extension and falling back to triton.language.tl.
• Codebase Update: Existing Triton kernels across several files (swiglu_quant.py, solve_tril.py, split_qkv_rmsnorm_rope.py, reject_sample.py) have been updated to use these newly unified functions.

🧠 New Feature in Public Preview: You can now enable Memory 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.

Changelog

• vllm_ascend/ops/triton/activation/swiglu_quant.py
  • Changed direct calls to tl.extract_slice to use the newly defined extract_slice utility function.
• vllm_ascend/ops/triton/fla/solve_tril.py
  • Changed direct calls to tl.insert_slice and tl.extract_slice to use the newly defined insert_slice and extract_slice utility functions.
• vllm_ascend/ops/triton/linearnorm/split_qkv_rmsnorm_rope.py
  • Changed direct calls to tl.extract_slice and tl.insert_slice to use the newly defined extract_slice and insert_slice utility functions.
• vllm_ascend/ops/triton/reject_sample.py
  • Changed direct calls to tl.get_element to use the newly defined get_element utility function.
• vllm_ascend/ops/triton/triton_utils.py
  • Introduced _resolve_triton_ascend_op to dynamically resolve Triton operations (insert_slice, extract_slice, get_element) from either triton.language.extra.cann.extension or triton.language.tl.
  • Exported insert_slice, extract_slice, and get_element as unified functions.

Activity

• No specific activity has been recorded for this pull request yet.

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 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

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 refactors the usage of Triton operators insert_slice, extract_slice, and get_element to improve compatibility with different versions of CANN. A new dispatch mechanism in vllm_ascend/ops/triton/triton_utils.py is introduced to resolve these operators, prioritizing versions from triton.language.extra.cann.extension before falling back to the standard triton.language. The rest of the changes apply this new unified dispatch across various Triton kernels.

The pull request description is currently incomplete, which violates the repository's style guide. To improve clarity and help with the review process, I've prepared a suggested title and summary that follow the required format.

Suggested PR Title:

[Triton][Refactor] Unify dispatch of slice and element ops for CANN compatibility

Suggested PR Summary:

### What this PR does / why we need it?

This pull request refactors the dispatch mechanism for the Triton operators `insert_slice`, `extract_slice`, and `get_element` to ensure compatibility with both CANN 8.5 and 9.0.

A unified helper function, `_resolve_triton_ascend_op`, has been introduced in `vllm_ascend/ops/triton/triton_utils.py`. This function dynamically resolves these operators by first attempting to import them from the `triton.language.extra.cann.extension` module, which is present in newer CANN versions. If that fails, it falls back to the standard `triton.language` module.

This approach centralizes operator dispatch logic, allowing individual Triton kernels to use these functions without being aware of the underlying Triton/CANN version. All call sites have been updated to use these new unified functions.

### Does this PR introduce _any_ user-facing change?

No. This is an internal refactoring of operator implementations and does not introduce any user-facing changes.

### How was this patch tested?

CI is expected to pass with existing tests.

**Testing Context:**
- vLLM version: v0.16.0
- vLLM main: `15d76f74e2fdb12a95ea00f0ca283acf6219a2b7`