ascend-migration-analysis

Ascend NPU Migration Analysis Skill

Purpose

Perform a systematic Ascend NPU migration feasibility analysis for any PyTorch project. Scan the codebase for CUDA/NVIDIA dependencies, classify them by domain, provide replacement solutions item by item, and estimate the migration effort.

When to Use

• The user asks whether a certain project/model can run on Ascend NPU
• The user needs to evaluate the effort of migrating a PyTorch project to Ascend NPU
• The user needs to conduct comprehensive technical due diligence for a CUDA→Ascend migration
• The user encounters NPU compatibility issues and needs to troubleshoot which CUDA dependencies have not been replaced

Analysis Workflow

Perform the following 7-step analysis on the target project, with each step corresponding to one migration domain. See the reference documents for each domain for detailed methods and replacement solutions.

Step 1: Third-Party Dependency Audit (Domain 1)

Read references/01-dependency-audit.md.

Scan contents:

• requirements.txt / setup.py / pyproject.toml NVIDIA-exclusive dependencies in
• flash-attn, triton, xformers, nvidia-ml-py, cuda-python, cupy , etc.
• CUDA hardcoding in indirect dependencies

Output:Dependency compatibility matrix(each dependency labeled: compatible / needs replacement / unsupported / not involved)

Step 2: Device Layer Scan (Domain 2)

Read references/02-device-layer.md.

Scan contents (using grep/rg):

• torch.cuda.* the entire series of APIs
• hardcoded "cuda" device strings
• backend="nccl" distributed backend
• amp.autocast(device_type='cuda') mixed precision
• torch.compile / @torch.compile usage
• torch.cuda.Stream / torch.cuda.Event
• init_device_mesh("cuda")

Output:Device layer replacement checklist(file:line number → replacement solution)

Step 3: Attention Mechanism Analysis (Domain 3)

Read references/03-attention-mechanism.md.

Scan contents:

• flash_attn (FA2) import and invocation
• flash_attn_interface (FA3) import and invocation
• xformers attention invocation
• F.scaled_dot_product_attention invocation
• manually implemented attention (Q×K^T softmax, etc.)
• attention tensor layout (BSHD / BNSD, etc.)

Output:Attention backend replacement solution(each attention type → corresponding NPU replacement)

Step 4: Custom Operator / Kernel Analysis (Domain 4)

Read references/04-custom-operators.md.

Scan contents:

• .cu / .cuh / .cpp custom CUDA kernel files
• @triton.jit Triton kernels decorated with
• torch.utils.cpp_extension.CUDAExtension build
• torch.autograd.Function custom forward/backward in subclasses
• csrc/ / ops/ / kernels/ directory

Output:Custom operator migration solution(functional description of each kernel + replacement path)

Step 5: Distributed Communication Analysis (Domain 5)

Read references/05-distributed.md.

Scan contents:

• dist.init_process_group backend selection
• dist.all_to_all / dist.all_gather / dist.broadcast , etc.
• init_device_mesh / DeviceMesh usage
• Ulysses / Ring / Sequence Parallel patterns
• torch.distributed.P2POp / batch_isend_irecv

Output:Distributed communication adaptation solution

Step 6: Precision & Numerics Analysis (Domain 6)

Read references/06-precision-strategy.md.

Scan contents:

• .float() / .double() type conversion
• torch.bfloat16 / torch.float16 dtype usage
• torch.complex128 / torch.float64 high-precision usage
• autocast dtype settings in context
• float32 protection in numerically sensitive operations (RMSNorm, LayerNorm, AdaLN)
• dtype == torch.float32 assert statements

Output:Precision strategy adjustment checklist

Step 7: Task-Phase Dependency Matrix (Domain 7)

Read references/07-task-phase-matrix.md.

Split the findings from the first 6 steps bytask execution stagesand identify:

• which dependencies are needed in every stage (P0)
• which dependencies are needed only in specific stages/modes (P1)
• Which dependencies are not involved in the current task (can be skipped)

Output:Dependency matrix by phase + Minimal migration set + Full migration effort estimate

Output Format

After the analysis is complete, output a structured report:

# {Project} Ascend NPU Migration Assessment

## Executive Summary
- Overall feasibility: [可行 / 需适配 / 困难 / 不可行]
- Estimated effort: [X 周]
- Blockers: [list]

## Dependency Matrix (per task phase)
| Dependency | Phase A | Phase B | Replacement | Effort |
|------------|---------|---------|-------------|--------|

## Detailed Findings (per domain)
### Domain 1: Third-Party Dependencies
### Domain 2: Device Layer
### Domain 3: Attention Mechanism
### Domain 4: Custom Operators
### Domain 5: Distributed Communication
### Domain 6: Precision Strategy

## Migration Roadmap
## Risk Matrix

Quick Reference: Common Replacements

CUDA Pattern	Ascend Replacement	Confidence
torch.cuda.empty_cache()	torch.npu.empty_cache() or transfer_to_npu	High
backend="nccl"	backend="hccl"	High
autocast('cuda')	autocast('npu')	High
device="cuda"	device="npu"	High
flash_attn.flash_attn_func	mindiesd.attention_forward(op_type="ascend_laser_attention")	High
flash_attn.flash_attn_varlen_func	mindiesd.attention_forward(op_type="fused_attn_score")	Medium
xformers.ops.memory_efficient_attention	mindiesd.attention_forward(op_type="fused_attn_score")	High
F.scaled_dot_product_attention	Natively supported by torch_npu	High
@triton.jit kernels	Needs to be analyzed individually; may be replaced by mindiesd/RainFusion/native NPU operators	Low-Medium
torch.compile	Disable or use the torch_npu backend	Medium
RMSNorm .float()	torch_npu.npu_rms_norm()	High
LayerNorm .float()	Remove .float(), NPU native BF16	High
RoPE complex128	Downgrade to complex64 or use mindiesd.rotary_position_embedding()	High
init_device_mesh("cuda")	"npu"	Medium
.cu / CUDAExtension	Needs to be rewritten as Ascend operators	Low

Notes

• This skill's alternative-solution knowledge base comes from actual Wan2.2 CUDA→Ascend migration experience
• Migration patterns for DiT/Transformer-type architectures are highly reusable
• Analysis results need to be verified against the actual CANN/torch_npu version
• Some alternative solutions depend on the mindiesd library; confirm whether it is available in the environment