torch-custom-ops-guide

Guides registration of custom PyTorch operators for npugraph_ex graph mode compilation.

Updated Sep 15, 2026
One-click install
npx skills add https://github.com/WangWindow/CANN-BatchMatMulMaxsum --skill torch-custom-ops-guide-wangwindow
Or copy as Structured Prompt for Agent
Please help me install this Agent Skill.
Skill: torch-custom-ops-guide
Source: https://github.com/WangWindow/CANN-BatchMatMulMaxsum/tree/main/.agents/skills/torch-custom-ops-guide
Command: npx skills add https://github.com/WangWindow/CANN-BatchMatMulMaxsum --skill torch-custom-ops-guide-wangwindow

SYSTEM DOCUMENTATION & REQUIREMENTS

What problem does it solve? Custom PyTorch operators that run in Eager mode cannot be captured by torch.compile or npugraph_ex graph mode without proper registration and Meta functions. This Skill walks you through the complete workflow of adapting custom operators for graph mode on Ascend NPU, covering both registration paths and Meta inference function authoring. ## Core Features & Use Cases - State-based routing: Identifies whether your operator is undeveloped, Eager-ready, or already in-graph, then routes to the correct workflow (template generation, adaptation, or triage). - Dual registration paths: Supports both torch.library.custom_op with register_fake and pure Python torch.library.Library with FRAGMENT/IMPL/Meta registration, including in-place (mutates_args, Tensor(a!)) versus out-of-place semantics. - Meta function authoring: Generates shape/dtype/device-only inference skeletons that enable graph capture without executing real computation. - Use Case: You have a FlashAttention operator running in Eager mode via torch.library.custom_op and want it compiled by npugraph_ex; the Skill guides you to write a register_fake Meta function in your user script so torch.compile can build the graph. ## Quick Start Ask the assistant how to adapt your existing Eager custom operator registered with torch.library.custom_op for npugraph_ex graph mode.

Frequently Asked Questions about torch-custom-ops-guide

High-intent search queries and answers about installing and using this skill.

FAQPage Schema
How do I adapt a custom PyTorch operator for torch.compile graph mode?

Register a Meta inference function for your operator so torch.compile can build the graph. With torch.library.custom_op, add @my_op.register_fake in your user script; with torch.library.Library, register a Meta implementation via Library(namespace, "IMPL", "Meta").

torch.library.custom_op vs torch.library.Library for custom operators?

torch.library.custom_op is a decorator-based approach suited for registering individual operators concisely. Pure Python torch.library.Library uses explicit define and impl calls, fitting batch registration or finer control over schema, Eager, and Meta implementations.

What does a register_fake Meta function do?

A register_fake Meta function derives output shape, dtype, and device from input metadata without running real computation. It enables torch.compile to construct the graph; in-place operators return None while out-of-place operators return an empty tensor with correct metadata.

Can I run the real NPU kernel inside a Meta function?

No, Meta functions must never execute real kernels or read actual data. They only infer output metadata from input metadata and return fake tensors; real computation belongs exclusively to the device implementation registered for PrivateUse1.

How do in-place operators differ in custom_op registration?

In-place operators declare mutates_args in torch.library.custom_op and use Tensor(a!) alias annotations in Library schemas. Their Meta implementations return None, whereas out-of-place operators return an empty tensor with the correct shape, dtype, and device.

What if my operator already enters the graph but fails at runtime?

Do not rewrite the operator registration; the adaptation steps are already complete. Route to the torch-npugraph-ex-dfx-triage workflow to diagnose whether the failure occurs during compilation, runtime execution, or performance stages.