Skip to content

[Draft] Add Ascend NPU optim Support#478

Draft
zhangtao0408 wants to merge 216 commits intovipshop:ascendfrom
zhangtao0408:dev
Draft

[Draft] Add Ascend NPU optim Support#478
zhangtao0408 wants to merge 216 commits intovipshop:ascendfrom
zhangtao0408:dev

Conversation

@zhangtao0408
Copy link
Contributor

@zhangtao0408 zhangtao0408 commented Nov 25, 2025

Optim: npu_fast_gelu, npu_rms_norm, npu_layer_norm_eval, npu_rotary_mul, npu_weight_nz, npu_adalayernorm

@DefTruth DefTruth self-requested a review November 25, 2025 07:07
@DefTruth
Copy link
Member

DefTruth commented Nov 25, 2025

Thanks! Please switch the target branch to ascend

@zhangtao0408 zhangtao0408 changed the base branch from dev to ascend November 25, 2025 07:10
@DefTruth DefTruth changed the title [Draft] Add Ascend NPU optim Support for Cache-Dit [Draft] Add Ascend NPU optim Support Nov 25, 2025
DefTruth and others added 26 commits November 26, 2025 18:36
@DefTruth
feat: support async ulysses cp for flux (vipshop#480)
2b00fce 
* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux

* feat: support async ulysses qkv proj for Flux
@DefTruth
chore: update support matrix (vipshop#484)
2cbb731
@DefTruth
chore: update async ulysses cp docs (vipshop#486)
1e3c69e
@DefTruth
chore: update async ulysses cp refs (vipshop#487)
e61498e
@gameofdimension
feat: support FLUX.2-dev Tensor Parallelism (vipshop#485)
92db44a 
* shard transformer

* tp text encoder

* refactor

* clean up

* save more gpu memory

---------

Co-authored-by: felix01.yu <felix01.yu@vipshop.com>
@DefTruth
feat: support Hybrid cache + TP for 🔥FLUX.2 (vipshop#489)
54461e5 
* misc: refactor flux.2 tensor parallel

* misc: refactor flux.2 tensor parallel

* misc: refactor flux.2 tensor parallel

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2

* feat: support hybrid cache + tp for flux.2
@DefTruth
feat: Add seq offload for 🔥FLUX.2 w/o parallel (vipshop#490)
4cae178 
* feat: support hybrid cache + tp for flux.2

* feat: enable seq offload for FLUX.2 w/ GPU=1
@DefTruth
feat: support 🔥FLUX.2 context parallel (vipshop#492)
b9fd68e 
* feat: support hybrid cache + tp for flux.2

* feat: enable seq offload for FLUX.2 w/ GPU=1

* feat: support FLUX.2 context parallel

* feat: support FLUX.2 context parallel

* feat: support FLUX.2 context parallel

* feat: support FLUX.2 context parallel

* feat: support FLUX.2 context parallel

* feat: support FLUX.2 context parallel
@BBuf @DefTruth
feat: support torch profiler in cache-dit (vipshop#491)
1f68bc1 
* opt run_wan_tp example

* upd

* support torch profiler in cache-dit

* upd

* upd

* upd

* upd

* upd

* upd

* Update README.md

---------

Co-authored-by: DefTruth <31974251+DefTruth@users.noreply.github.com>
@gameofdimension
feat: support 🔥z-image tensor parallel (vipshop#494)
bae27d6 
* add zimage

* inspect

* apply tp plan

* apply tp plan

* fix

* fix bug

---------

Co-authored-by: felix01.yu <felix01.yu@vipshop.com>
@gameofdimension
feat: support lumina2 tensor parallel (vipshop#495)
245c8ab 
* add lumina tp

* fix name

* update

---------

Co-authored-by: felix01.yu <felix01.yu@vipshop.com>
@DefTruth
feat: support cache for 🔥z-image (vipshop#496)
bd2f685 
* feat: support cache for z-image

* feat: support cache for z-image

* feat: support cache for z-image
@DefTruth
feat: support context parallel for 🔥z-image (vipshop#497)
4abc829 
* feat: support context parallel for z-image

* feat: support context parallel for z-image

* feat: support context parallel for z-image

* feat: support context parallel for z-image

* feat: support context parallel for z-image

* feat: support context parallel for z-image

* feat: support context parallel for z-image

* feat: support context parallel for z-image

* feat: support context parallel for z-image
@DefTruth
fix: temp FnB(n>0) workaround for z-image cache w/ cp (vipshop#499)
73fbd44 
* feat: support context parallel for z-image

* feat: support context parallel for z-image
@BBuf
add profiler for flux tp and cp example (vipshop#501)
29a8345
@DefTruth
chore: Update README.md (vipshop#502)
879e85c 
* chore: Update README.md

* Update README.md

* Update README.md

* Update README.md
@DefTruth
feat: support FnB0 for z-image w/ cp (vipshop#503)
9a5cdd9 
* feat: support FnB0 for z-image w/ cp

* feat: support FnB0 for z-image w/ cp

* feat: support FnB0 for z-image w/ cp
@DefTruth
feat: support _sdpa_cudnn backend for cp (vipshop#504)
8bd89ba 
* feat: fast rope for z-image

* chore: update notes

* chore: update z-image cp example

* chore: update z-image cp example

* chore: update z-image cp example

* chore: update z-image cp example

* chore: update z-image cp example

* feat: allow cudnn attn w/ attn mask for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp

* feat: support _sdpa_cudnn backend for cp
@DefTruth
feat: support async ulysses cp for z-image (vipshop#505)
e0aabd7 
* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image
@DefTruth
feat: add all_to_all_single v2 (vipshop#507)
0f891ea 
* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image

* feat: support async ulysses cp for z-image

* feat: add all_to_all_single v2

* feat: add all_to_all_single v2
@DefTruth
feat: support async ulysses cp for qwen-image (vipshop#508)
ed9a226 
* feat: support async ulysses cp for qwen-image

* feat: support async ulysses cp for qwen-image

* feat: support async ulysses cp for qwen-image
@triple-mu @DefTruth
feat: support all2all qkv per token fp8 (vipshop#509)
6137662 
* support all2all_qkv fp8

* support qkv fp8 all2all

* support o fp8 all2all

* add wait_tensor

* add addition ulysses_anything_float8 logic

* Update utils.py

* fix log info error

* Update __init__.py

* Update _templated_ulysses_anything.py

---------

Co-authored-by: Your Name <you@example.com>
Co-authored-by: DefTruth <31974251+DefTruth@users.noreply.github.com>
@BBuf
chore: improve flux2 and qwen examples (vipshop#512)
ad1db02 
* add profiler for flux tp and cp example

* Enhance Flux2 & Qwen examples with customizable CLI arguments
@triple-mu @DefTruth
fix uaa-fp8 .view compile error (vipshop#514)
88d7ae0 
* uaa-fp8 support torch.compile

* Update _distributed_primitives.py

* Update _distributed_primitives.py

---------

Co-authored-by: DefTruth <31974251+DefTruth@users.noreply.github.com>
@DefTruth
feat: relaxed transformer strict assert (vipshop#515)
60b3f2f 
* feat: relaxed assert transformer

* feat: relaxed assert transformer

* feat: relaxed assert transformer
@DefTruth
feat: all2all qkv fp8 for ulysses (vipshop#516)
1dd51a2 
* feat: all2all qkv fp8 for ulysses

* feat: all2all qkv fp8 for ulysses

* feat: all2all qkv fp8 for ulysses

* feat: all2all qkv fp8 for ulysses

* feat: all2all qkv fp8 for ulysses
DefTruth and others added 15 commits January 15, 2026 16:54
@DefTruth
fix kwargs override error in load configs (vipshop#708)
f40e2ab
@DefTruth
CI: use torch-cpu for basic cpu tests (vipshop#709)
d808fff 
* CI: use torch-cpu for basic cpu tests

* CI: use torch-cpu for basic cpu tests

* CI: use torch-cpu for basic cpu tests
@DefTruth
docs: update vllm-omni x cache-dit link (vipshop#710)
cd60960
@DefTruth
chore: add flux fill examples (vipshop#711)
c7398ea
@luren55
npu docs update (vipshop#712)
f474da9 
* npu docs update

* npu docs update
@luren55
npu docs update (vipshop#713)
ce429ae 
* npu docs update

* npu docs update

* npu docs update
@luren55
NPU support update (vipshop#715)
09bac0a 
* npu docs update

* npu docs update

* npu docs update

* NPU support update
@gameofdimension
fix dim computation for FLUX.2 klein (vipshop#716)
74d80b2
@DefTruth @gameofdimension
feat: support 🚀flux2-klein series (vipshop#717)
128069f 
* feat: support flux2-klein

* feat: support flux2-klein

* feat: support flux2-klein

* feat: support flux2-klein

* feat: support flux2-klein series

* feat: support flux2-klein series

* feat: support flux2-klein series

* fix dim computation for FLUX.2 klein

* feat: support flux2-klein series

* feat: support flux2-klein series

* feat: support flux2-klein series

---------

Co-authored-by: G.O.D <32255912+gameofdimension@users.noreply.github.com>
@DefTruth
chore: use new logo (vipshop#718)
dfd4b2b 
* chore: use new logo

* chore: use new logo

* chore: use new logo

* chore: use new logo

* chore: use new logo

* chore: use new logo

* chore: use new logo

* chore: use new logo

* chore: use new logo
@DefTruth
fix logo link (vipshop#719)
be1bb83
@DefTruth
chore: enable scm for examples (vipshop#720)
30ba424 
* fix logo link

* enable scm for examples

* enable scm for examples
@DefTruth
chore: use new logo in mkdocs.yml (vipshop#721)
8c047c3
@DefTruth
docs: update installation user guide (vipshop#722)
853df89
@BBuf
[chore] Align torch generator with example (vipshop#723)
a75920d 
* ud

* ud
BBuf and others added 14 commits January 19, 2026 09:44
@BBuf @DefTruth
Fix generator bug in cache-dit (vipshop#724)
94b4e16 
* random

* Update bench.py

* Update bench_distill.py

---------

Co-authored-by: DefTruth <31974251+DefTruth@users.noreply.github.com>
@DefTruth
examples: allow custom generator device (vipshop#726)
00a2264 
* chore: allow custom generator device in examples

* chore: allow custom generator device in examples

* chore: allow custom generator device in examples

* chore: allow custom generator device in examples

* fix docs

* fix docs

* fix docs

* fix docs
@DefTruth
examples: allow custom warmup-steps (vipshop#727)
95b32ff
@DefTruth
docs: add latest news (vipshop#728)
410cb87 
* docs: add latest news

* docs: add latest news

* docs: add latest news

* docs: add latest news
@DefTruth
docs: fix docs format (vipshop#729)
1585482 
* docs: fix docs format

* docs: fix docs format
@66RING
fix selected metrics print (vipshop#730)
c4c0d82
@DefTruth
docs: add flux examples to tp docs (vipshop#731)
328edea
@DefTruth
fix ltx-2 i2v example (vipshop#734)
7c26b25 
* fix ltx-2 i2v example

* fix ltx-2 i2v example
@DefTruth
Update README.md (vipshop#735)
d130bc8 
* Update README.md

* Update README.md

* Update README.md
@DefTruth
chore: allow use default steps for scm (vipshop#736)
a677008 
* allow use default steps for scm

* allow use default steps for scm

* allow use default steps for scm

* allow use default steps for scm

* allow use default steps for scm

* allow use default steps for scm

* allow use default steps for scm
@BBuf
chore: support gpu generator in server (vipshop#737)
4fd633b
@DefTruth
docs: update download badge (vipshop#738)
38de92b
@BBuf
docs: refine profiler and serving docs (vipshop#739)
e887a55 
* refine docs

* add pert

* format
@DefTruth
fix UAA broken while using joint attn
bb033b3
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@DefTruth DefTruth Awaiting requested review from DefTruth

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.