[PR #154] [CLOSED] [Feat][Ascend] Add Ascend RoPE and RMS-RoPE backend #182

New issue

Closed

opened 2026-10-26 22:21:19 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:21:19 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/154
Author: @Big2Wheel
Created: 9/4/2026
Status: ❌ Closed

Base: main ← Head: codex/ascend-rope-rms-rope

📝 Commits (6)

5942fae feat(npu): add Ascend quantization backend
1ecb943 test(npu): include NPU in supported devices
0bccee7 refactor(ascend): clarify backend naming
5018b5e feat(ascend): add RoPE and RMS-RoPE backend
436dcb7 fix(ascend): validate required operator schemas
4d5513a fix(ascend): decouple backend capabilities

📊 Changes

9 files changed (+1248 additions, -41 deletions)

View changed files

📝 README.md (+63 -38)
📝 comfy_kitchen/__init__.py (+1 -0)
➕ comfy_kitchen/backends/ascend/__init__.py (+379 -0)
➕ comfy_kitchen/backends/ascend/rope.py (+365 -0)
📝 comfy_kitchen/registry.py (+1 -1)
📝 tests/conftest.py (+2 -2)
➕ tests/test_ascend_backend.py (+137 -0)
➕ tests/test_ascend_rope.py (+292 -0)
📝 tests/test_backends.py (+8 -0)

📄 Description

What this PR does / why we need it

Adds an optional Huawei Ascend backend for Comfy Kitchen's RoPE and RMS-RoPE APIs.

Supported public APIs:

apply_rope, apply_rope1, apply_rope_split_half, and apply_rope_split_half1
rms_rope, rms_rope1, rms_rope_split_half, and rms_rope_split_half1
The corresponding in-place variants of all APIs above

The backend uses torch_npu.npu_rotary_mul for rotary embeddings and torch_npu.npu_rms_norm for fused RMS normalization. It supports FP16/BF16, BNSD/BSND layouts, interleaved and split-half rotary conventions, packed-QKV strided views with a contiguous final dimension, GQA-style head counts, and partial rotary dimensions where exposed by the public API.

Capability checks reject unsupported shapes, layouts, dtypes, and broadcast patterns before dispatch, allowing the registry to select another capable backend. Runtime operator errors are not silently swallowed.

Relationship to #151 and #153

This is a stacked PR that currently depends on #151, which introduces the generic Ascend backend registration, availability detection, and baseline quantization support. The PR-specific commits to review on top of #151 are the RoPE implementation and its two capability-validation follow-ups.

#153 is a sibling feature PR, not a dependency of this PR. Both #153 and this PR branch directly from #151:

#151 Ascend backend foundation
├── #153 Ascend INT8 linear
└── this PR: Ascend RoPE / RMS-RoPE

After #151 merges, this branch will be rebased onto main; the PR diff will then contain only the RoPE/RMS-RoPE changes. No INT8-linear changes from #153 are included here.

Does this PR introduce any user-facing change?

Yes. On supported Ascend NPU installations, the registry can select the new Ascend implementation for compatible RoPE and RMS-RoPE calls. Other devices and unsupported Ascend call patterns retain the existing backend fallback behavior.

How was this patch tested?

Hardware: Huawei Ascend 910B3, restricted to one NPU with ASCEND_RT_VISIBLE_DEVICES=7.

Tests

pytest -q tests/test_ascend_backend.py tests/test_ascend_rope.py tests/test_backends.py\n109 passed after rebasing onto comfy-kitchen v0.2.33

pytest -q tests/test_ascend_backend.py tests/test_ascend_rope.py \
  tests/test_rope.py tests/test_rms_rope.py tests/test_backends.py \
  tests/test_constraints.py
319 passed, 404 skipped

Coverage includes FP16/BF16, BSND/BNSD, pair/single APIs, interleaved/split-half conventions, in-place/out-of-place behavior, packed QKV views, partial rotary dimensions, GQA head counts, input preservation, alias rejection, unsupported-call fallback, and automatic backend selection.

ruff check and git diff --check also pass for all changed Python files.

Profiler verification

A torch profiler trace confirms that the implementation reaches the Ascend operators:

npu::npu_rms_norm
aclnnRmsNorm
npu::npu_rotary_mul
aclnnRotaryPositionEmbeddingV2

No layout-induced contiguous operation appears in the profiled supported path.

Operator benchmark

Each measurement used 10 warm-up iterations and 30 synchronized measured iterations. Times include coefficient conversion and required dtype semantics.

Workload shape	Operation	Eager mean	Ascend mean	Improvement
Z-Image `[1, 4096, 30, 128]`	RoPE	1.206 ms	0.551 ms	54.2%
Z-Image `[1, 4096, 30, 128]`	RMS-RoPE	1.404 ms	0.638 ms	54.5%
Wan `[1, 12288, 16, 128]`	RoPE	4.023 ms	0.614 ms	84.8%
Wan `[1, 12288, 16, 128]`	RMS-RoPE	4.505 ms	0.614 ms	86.3%

End-to-end workflow validation

Z-Image BF16 showed the expected operator dispatch (272 RMS-RoPE calls per request), but no stable end-to-end improvement, so no image-workflow speedup is claimed.

For a fixed-seed Wan2.2 14B text-to-video workflow at 640x640 and 81 frames (320 RoPE calls per request):

Version	Hot-run mean	Median	Std. dev.
#151 baseline	96.34 s	95.97 s	1.44 s
This PR	91.79 s	91.71 s	0.39 s

This is a 4.73% mean end-to-end reduction. Decoding both generated videos and comparing all 81 frames produced zero differing pixel values (max_diff = 0).

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-kitchen/pull/154 **Author:** [@Big2Wheel](https://github.com/Big2Wheel) **Created:** 9/4/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `codex/ascend-rope-rms-rope` --- ### 📝 Commits (6) - [`5942fae`](https://github.com/Comfy-Org/comfy-kitchen/commit/5942fae7a80bc3c5df640b661c2fbd1d2b5fb092) feat(npu): add Ascend quantization backend - [`1ecb943`](https://github.com/Comfy-Org/comfy-kitchen/commit/1ecb9434690dc771152d108038da8612ff9f1ff4) test(npu): include NPU in supported devices - [`0bccee7`](https://github.com/Comfy-Org/comfy-kitchen/commit/0bccee7393739704ae3c05ab233908820df51eef) refactor(ascend): clarify backend naming - [`5018b5e`](https://github.com/Comfy-Org/comfy-kitchen/commit/5018b5eadc66471284122d413990180321759809) feat(ascend): add RoPE and RMS-RoPE backend - [`436dcb7`](https://github.com/Comfy-Org/comfy-kitchen/commit/436dcb7760cdd871b96f6ab5a19467082ee080a4) fix(ascend): validate required operator schemas - [`4d5513a`](https://github.com/Comfy-Org/comfy-kitchen/commit/4d5513a5ffb91263514b2d9ad90bccaece943b08) fix(ascend): decouple backend capabilities ### 📊 Changes **9 files changed** (+1248 additions, -41 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+63 -38) 📝 `comfy_kitchen/__init__.py` (+1 -0) ➕ `comfy_kitchen/backends/ascend/__init__.py` (+379 -0) ➕ `comfy_kitchen/backends/ascend/rope.py` (+365 -0) 📝 `comfy_kitchen/registry.py` (+1 -1) 📝 `tests/conftest.py` (+2 -2) ➕ `tests/test_ascend_backend.py` (+137 -0) ➕ `tests/test_ascend_rope.py` (+292 -0) 📝 `tests/test_backends.py` (+8 -0) </details> ### 📄 Description ## What this PR does / why we need it Adds an optional Huawei Ascend backend for Comfy Kitchen's RoPE and RMS-RoPE APIs. Supported public APIs: - `apply_rope`, `apply_rope1`, `apply_rope_split_half`, and `apply_rope_split_half1` - `rms_rope`, `rms_rope1`, `rms_rope_split_half`, and `rms_rope_split_half1` - The corresponding in-place variants of all APIs above The backend uses `torch_npu.npu_rotary_mul` for rotary embeddings and `torch_npu.npu_rms_norm` for fused RMS normalization. It supports FP16/BF16, BNSD/BSND layouts, interleaved and split-half rotary conventions, packed-QKV strided views with a contiguous final dimension, GQA-style head counts, and partial rotary dimensions where exposed by the public API. Capability checks reject unsupported shapes, layouts, dtypes, and broadcast patterns before dispatch, allowing the registry to select another capable backend. Runtime operator errors are not silently swallowed. ## Relationship to #151 and #153 This is a stacked PR that currently depends on #151, which introduces the generic Ascend backend registration, availability detection, and baseline quantization support. The PR-specific commits to review on top of #151 are the RoPE implementation and its two capability-validation follow-ups. #153 is a sibling feature PR, not a dependency of this PR. Both #153 and this PR branch directly from #151: ```text #151 Ascend backend foundation ├── #153 Ascend INT8 linear └── this PR: Ascend RoPE / RMS-RoPE ``` After #151 merges, this branch will be rebased onto `main`; the PR diff will then contain only the RoPE/RMS-RoPE changes. No INT8-linear changes from #153 are included here. ## Does this PR introduce any user-facing change? Yes. On supported Ascend NPU installations, the registry can select the new Ascend implementation for compatible RoPE and RMS-RoPE calls. Other devices and unsupported Ascend call patterns retain the existing backend fallback behavior. ## How was this patch tested? Hardware: Huawei Ascend 910B3, restricted to one NPU with `ASCEND_RT_VISIBLE_DEVICES=7`. ### Tests ```text pytest -q tests/test_ascend_backend.py tests/test_ascend_rope.py tests/test_backends.py\n109 passed after rebasing onto comfy-kitchen v0.2.33 pytest -q tests/test_ascend_backend.py tests/test_ascend_rope.py \ tests/test_rope.py tests/test_rms_rope.py tests/test_backends.py \ tests/test_constraints.py 319 passed, 404 skipped ``` Coverage includes FP16/BF16, BSND/BNSD, pair/single APIs, interleaved/split-half conventions, in-place/out-of-place behavior, packed QKV views, partial rotary dimensions, GQA head counts, input preservation, alias rejection, unsupported-call fallback, and automatic backend selection. `ruff check` and `git diff --check` also pass for all changed Python files. ### Profiler verification A torch profiler trace confirms that the implementation reaches the Ascend operators: ```text npu::npu_rms_norm aclnnRmsNorm npu::npu_rotary_mul aclnnRotaryPositionEmbeddingV2 ``` No layout-induced `contiguous` operation appears in the profiled supported path. ### Operator benchmark Each measurement used 10 warm-up iterations and 30 synchronized measured iterations. Times include coefficient conversion and required dtype semantics. | Workload shape | Operation | Eager mean | Ascend mean | Improvement | | --- | --- | ---: | ---: | ---: | | Z-Image `[1, 4096, 30, 128]` | RoPE | 1.206 ms | 0.551 ms | 54.2% | | Z-Image `[1, 4096, 30, 128]` | RMS-RoPE | 1.404 ms | 0.638 ms | 54.5% | | Wan `[1, 12288, 16, 128]` | RoPE | 4.023 ms | 0.614 ms | 84.8% | | Wan `[1, 12288, 16, 128]` | RMS-RoPE | 4.505 ms | 0.614 ms | 86.3% | ### End-to-end workflow validation Z-Image BF16 showed the expected operator dispatch (272 RMS-RoPE calls per request), but no stable end-to-end improvement, so no image-workflow speedup is claimed. For a fixed-seed Wan2.2 14B text-to-video workflow at 640x640 and 81 frames (320 RoPE calls per request): | Version | Hot-run mean | Median | Std. dev. | | --- | ---: | ---: | ---: | | #151 baseline | 96.34 s | 95.97 s | 1.44 s | | This PR | 91.79 s | 91.71 s | 0.39 s | This is a 4.73% mean end-to-end reduction. Decoding both generated videos and comparing all 81 frames produced zero differing pixel values (`max_diff = 0`). --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 22:21:19 +00:00