[PR #155] [CLOSED] [Feat][Ascend] Add fused RotateQuant backend #181

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/155
Author: @Big2Wheel
Created: 9/4/2026
Status: ❌ Closed

Base: main ← Head: codex/ascend-rotate-quant

📝 Commits (7)

5942fae feat(npu): add Ascend quantization backend
1ecb943 test(npu): include NPU in supported devices
0bccee7 refactor(ascend): clarify backend naming
6ab910a feat(ascend): add INT8 linear backend
91a5a3f test(ascend): handle optional quant matmul
8651a6d feat(ascend): add fused rotate quant backend
d8dd6d8 fix(ascend): handle zero quantization scales

📊 Changes

7 files changed (+956 additions, -41 deletions)

View changed files

📝 README.md (+66 -38)
📝 comfy_kitchen/__init__.py (+1 -0)
➕ comfy_kitchen/backends/ascend/__init__.py (+445 -0)
📝 comfy_kitchen/registry.py (+1 -1)
📝 tests/conftest.py (+2 -2)
➕ tests/test_ascend_backend.py (+424 -0)
📝 tests/test_backends.py (+17 -0)

📄 Description

What this PR does / why we need it

Adds an optional Huawei Ascend implementation of quantize_and_rotate_rowwise using torch_npu.npu_rotate_quant and integrates it into the Ascend ConvRot INT8 linear path.

The backend:

registers the capability only when npu_rotate_quant is available;
validates dtype, shape, group size, device, and deterministic-rounding requirements;
supports FP16/BF16, batched inputs, and non-contiguous inputs;
falls back to the existing _rotate_activation + npu_dynamic_quant path for unsupported runtimes or shapes;
avoids an extra scale-clamp kernel in the internal INT8 linear hot path while preserving the public API's positive-scale contract.

Relationship to the other Ascend PRs

Depends on #151, which introduces the base Ascend backend and quantization capabilities.
Depends on #153, which introduces the Ascend INT8 linear and ConvRot integration point. This branch is stacked directly on the latest #153 branch.
Does not depend on #154. The RoPE/RMS-RoPE backend is orthogonal and can be reviewed or merged independently.

Once #151 and #153 are merged, this PR can be rebased onto main, leaving only the RotateQuant-specific commits.

Performance and memory

Measured on one Ascend 910B3 with real Z-Image activation shapes:

Shape	Eager API	Ascend RotateQuant	Reduction
`[32, 3840]`	0.487 ms	0.174 ms	64.3%
`[8192, 3840]`	1.337 ms	0.207 ms	84.6%
`[32, 10240]`	0.466 ms	0.169 ms	63.7%
`[8192, 10240]`	2.945 ms	0.457 ms	84.5%

NPU profiling at [8192, 3840] confirms that aclnnMatmul + aclnnDynamicQuantV2 is replaced by aclnnRotateQuant, reducing the combined device time by 18.05% across ten iterations.

An alternating AB/BA Z-Image INT8 workflow test found no statistically meaningful end-to-end latency change; generated PNGs were pixel-identical. The fused operator uses additional temporary workspace: approximately 60 MiB at [8192, 3840] and 160 MiB at [8192, 10240] relative to the separate path.

Does this PR introduce any user-facing change?

Yes. Eligible quantize_and_rotate_rowwise calls and ConvRot Ascend INT8 linear calls automatically use the fused NPU operator. Unsupported configurations retain the existing fallback behavior.

How was this patch tested?

Ascend 910B3 operator correctness against the eager reference for FP16 and BF16.
Group sizes 16, 64, and 256.
Batched, all-zero-row, non-contiguous, and asymmetric-rotation inputs.
Capability selection, unsupported-shape decline, missing-operator fallback, and INT8 linear integration tests.
Caller-provided zero tensor-wise scales for FP16 and BF16, matching eager quantization semantics.
Fixed-seed Z-Image INT8 workflow: baseline and fused outputs were pixel-identical.
Wan2.2 14B BF16 smoke test completed successfully and made zero npu_rotate_quant calls.
Focused tests: 55 passed.
Full post-rebase suite: 660 passed, 1603 skipped, 1 warning.
ruff check, ruff format --check, and git diff --check passed.

_{🔄 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/155 **Author:** [@Big2Wheel](https://github.com/Big2Wheel) **Created:** 9/4/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `codex/ascend-rotate-quant` --- ### 📝 Commits (7) - [`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 - [`6ab910a`](https://github.com/Comfy-Org/comfy-kitchen/commit/6ab910aad706e9ee28d140f36e964649165e909c) feat(ascend): add INT8 linear backend - [`91a5a3f`](https://github.com/Comfy-Org/comfy-kitchen/commit/91a5a3fed39a8b88858aaa0b9e214a59ccc32e90) test(ascend): handle optional quant matmul - [`8651a6d`](https://github.com/Comfy-Org/comfy-kitchen/commit/8651a6dfe8a9cbaa3cc9eac1a8ab05b9f0d8584c) feat(ascend): add fused rotate quant backend - [`d8dd6d8`](https://github.com/Comfy-Org/comfy-kitchen/commit/d8dd6d8d2b0a5d1bd32ca687a69ab782d528782c) fix(ascend): handle zero quantization scales ### 📊 Changes **7 files changed** (+956 additions, -41 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+66 -38) 📝 `comfy_kitchen/__init__.py` (+1 -0) ➕ `comfy_kitchen/backends/ascend/__init__.py` (+445 -0) 📝 `comfy_kitchen/registry.py` (+1 -1) 📝 `tests/conftest.py` (+2 -2) ➕ `tests/test_ascend_backend.py` (+424 -0) 📝 `tests/test_backends.py` (+17 -0) </details> ### 📄 Description ## What this PR does / why we need it Adds an optional Huawei Ascend implementation of `quantize_and_rotate_rowwise` using `torch_npu.npu_rotate_quant` and integrates it into the Ascend ConvRot INT8 linear path. The backend: - registers the capability only when `npu_rotate_quant` is available; - validates dtype, shape, group size, device, and deterministic-rounding requirements; - supports FP16/BF16, batched inputs, and non-contiguous inputs; - falls back to the existing `_rotate_activation` + `npu_dynamic_quant` path for unsupported runtimes or shapes; - avoids an extra scale-clamp kernel in the internal INT8 linear hot path while preserving the public API's positive-scale contract. ### Relationship to the other Ascend PRs - Depends on #151, which introduces the base Ascend backend and quantization capabilities. - Depends on #153, which introduces the Ascend INT8 linear and ConvRot integration point. This branch is stacked directly on the latest #153 branch. - Does not depend on #154. The RoPE/RMS-RoPE backend is orthogonal and can be reviewed or merged independently. Once #151 and #153 are merged, this PR can be rebased onto `main`, leaving only the RotateQuant-specific commits. ## Performance and memory Measured on one Ascend 910B3 with real Z-Image activation shapes: | Shape | Eager API | Ascend RotateQuant | Reduction | | --- | ---: | ---: | ---: | | `[32, 3840]` | 0.487 ms | 0.174 ms | 64.3% | | `[8192, 3840]` | 1.337 ms | 0.207 ms | 84.6% | | `[32, 10240]` | 0.466 ms | 0.169 ms | 63.7% | | `[8192, 10240]` | 2.945 ms | 0.457 ms | 84.5% | NPU profiling at `[8192, 3840]` confirms that `aclnnMatmul + aclnnDynamicQuantV2` is replaced by `aclnnRotateQuant`, reducing the combined device time by 18.05% across ten iterations. An alternating AB/BA Z-Image INT8 workflow test found no statistically meaningful end-to-end latency change; generated PNGs were pixel-identical. The fused operator uses additional temporary workspace: approximately 60 MiB at `[8192, 3840]` and 160 MiB at `[8192, 10240]` relative to the separate path. ## Does this PR introduce any user-facing change? Yes. Eligible `quantize_and_rotate_rowwise` calls and ConvRot Ascend INT8 linear calls automatically use the fused NPU operator. Unsupported configurations retain the existing fallback behavior. ## How was this patch tested? - Ascend 910B3 operator correctness against the eager reference for FP16 and BF16. - Group sizes 16, 64, and 256. - Batched, all-zero-row, non-contiguous, and asymmetric-rotation inputs. - Capability selection, unsupported-shape decline, missing-operator fallback, and INT8 linear integration tests. - Caller-provided zero tensor-wise scales for FP16 and BF16, matching eager quantization semantics. - Fixed-seed Z-Image INT8 workflow: baseline and fused outputs were pixel-identical. - Wan2.2 14B BF16 smoke test completed successfully and made zero `npu_rotate_quant` calls. - Focused tests: `55 passed`. - Full post-rebase suite: `660 passed, 1603 skipped, 1 warning`. - `ruff check`, `ruff format --check`, and `git diff --check` passed. --- <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