[core] Support capture custom ops into aclgraph

[ganyi1996ppo]

What this PR does / why we need it?

Thanks to the PR #426 make vllm-ascend support the aclgraph inference to reduce the host overhead. However, the capability of aclgraph strongly relies on the functionality provided by torch.compile, which is the key feature supported in torch 2.x . Therefore, capture custom op into aclgraph is only possible when it can be recognize and captured by torch.compile.

In this PR, we register the meta implementation of current custom ops to enable the fx graph capture. And by doing that, insert those custom ops into aclgraph become a natural thing to the ascend runtime.

Does this PR introduce any user-facing change?

No user face change.

How was this patch tested?

Tested in unittest, we will integrate the rotary_embedding op into a small custom model and use torch.compile and aclgraph to capture and replay it to verify its functionality.

• vLLM version: v0.10.0
• vLLM main: vllm-project/vllm@099c046

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

[ganyi1996ppo]

@ttanzhiqiang Please review this PR also

[ganyi1996ppo]

@yiz-liu Please review this PR also

[codecov]

Codecov Report

❌ Patch coverage is 51.85185% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.04%. Comparing base (0bd5ff5) to head (951506e).

Files with missing lines	Patch %	Lines
vllm_ascend/meta_registration.py	50.00%	13 Missing ⚠️

❌ Your patch check has failed because the patch coverage (51.85%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2113      +/-   ##
==========================================
- Coverage   76.09%   76.04%   -0.05%     
==========================================
  Files         114      115       +1     
  Lines       13103    13130      +27     
==========================================
+ Hits         9971     9985      +14     
- Misses       3132     3145      +13     

Flag	Coverage Δ
unittests	76.04% <51.85%> (-0.05%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ttanzhiqiang]

This place is a static shape. If the shape of static_positions, static_hidden_states has changed, does meta need to go through again?

[ganyi1996ppo]

yes

[ttanzhiqiang]

Can custom ops be specially written to handle the problem of aclgraph without adding an operator to a meta at a time

[ganyi1996ppo]

Can custom ops be specially written to handle the problem of aclgraph without adding an operator to a meta at a time

No for now. We can try to write a macro or template to automatically generate the meta implementation, but no plan for this yet. We can consider this after we have enough number of custom ops.

[wangxiyuan]

LGTM, Please make the CI happy, I think we should merge this in high priority

[ganyi1996ppo]

LGTM, Please make the CI happy, I think we should merge this in high priority

The aclgraph path seems have unaligned accuracy compared with eager path on CI, I can't reproduce it on my local environment.

[MengqingCao]

It seems the first failed case is acc test case due to the model download issue. I think you can rebase your code, as acc test will not run per pr now.

[ganyi1996ppo]

It seems the first failed case is acc test case due to the model download issue. I think you can rebase your code, as acc test will not run per pr now.

Got, I'll give it a try