[PR #161] [CLOSED] [Feat][Ascend] Add ConvRot W4A4 linear backend #186

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: codex/ascend-convrot-w4a4

📝 Commits (10+)

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
9ad6f5e feat(ascend): add ConvRot W4A4 linear backend
aaad462 fix(ascend): guard RotateQuant feature limit
8232c93 test(ascend): cover fallback without RotateQuant

📊 Changes

7 files changed (+1368 additions, -45 deletions)

View changed files

📝 README.md (+76 -38)
📝 comfy_kitchen/__init__.py (+1 -0)
➕ comfy_kitchen/backends/ascend/__init__.py (+595 -0)
📝 comfy_kitchen/registry.py (+1 -1)
📝 tests/conftest.py (+2 -2)
➕ tests/test_ascend_backend.py (+668 -0)
📝 tests/test_backends.py (+25 -4)

📄 Description

What this PR does / why we need it

Adds an optional Ascend implementation of convrot_w4a4_linear, prioritizing the eager reference's numerical behavior.

Correctness revision

The previous implementation silently converted FP32 activations to BF16 before rotation/quantization, and narrowed all matmul results through FP16. An audit demonstrated changed INT4 activation codes, fixed-seed image divergence, and finite FP32/BF16 inputs producing non-finite outputs where eager remained finite.

This revision:

preserves the input dtype during rotation and reuses the eager signed-INT4 quantizer (including the [-7,7] emission range);
unpacks the original A4/W4 integer codes to INT8 and calls npu_quant_matmul with a unit FLOAT scale, no per-token scale, and INT32 output;
applies the cast, activation scale, weight scale, and bias in eager's order, with no compulsory FP16 intermediate for FP32/BF16;
rejects feature widths that could overflow the INT32 accumulator;
no longer requires RotateQuant for this capability; existing INT8/RotateQuant APIs and their width guards are unchanged;
retains fallback through the registry for unsupported call contracts.

This is not a zero-copy packed-INT4 matmul implementation. It trades unpacking and reference-preprocessing cost for correctness. FP16 retains eager's own intermediate-range limitations; the implementation does not promise greater precision than the reference.

Correction of earlier claims

Both eager and Ascend paths quantize activations to A4. The earlier explanation that the image difference was expected simply because the optimized path quantizes activations was incorrect.

The previously reported 134.00 s / 2.73x result belongs to the superseded, numerically different implementation. It is withdrawn as a performance claim for the current revision. New diagnostic timings are not a repeated hot A/B benchmark and no speedup is claimed here.

Relationship to the other Ascend PRs

#151 supplies the base optional Ascend backend/registration.
#153 supplies the quantized-matmul backend infrastructure in the same module. This implementation does not call the public int8_linear API or change its semantics.
This branch remains stacked on #155 for review/history, but W4A4 no longer calls its fused RotateQuant helper. The Git stack dependency is distinct from the runtime requirement.
#154 (RoPE/RMS-RoPE) is orthogonal and is not included in this stack.

After the predecessor stack lands, this PR can be rebased on main to leave the W4A4-specific changes.

The branch also merges main through 21003fa, retaining both the upstream allocation-context import and optional Ascend registration to resolve their import conflict. The same regression suite passed again after this merge.

Does this PR introduce any user-facing change?

Yes. Eligible Ascend ConvRot W4A4 calls use the correctness-first integer-accumulation implementation. Weights remain INT4-quantized on disk/in the caller, but the backend materializes INT8 values for matmul. No new mandatory dependency or opt-in environment variable is introduced.

How was this patch tested?

Hardware/runtime: one Ascend 910B3 (physical card 4), PyTorch 2.10.0+cpu with torch_npu 2.10.0.post4, CANN 9.1.0.

Zero-tolerance comparisons against eager for FP32, BF16 and FP16, including K=6144 and K=16384.
Large finite FP32/BF16 outputs beyond the FP16 range, zero/tiny/non-finite rows, bias, batched strided inputs, and signed/unsigned packed storage.
Operator-spy verification of original quantized codes, INT8 inputs, INT32 output, unit scale, no per-token scale, and no weight mutation.
Independent CPU INT64 oracle for the native integer accumulator, including negative INT4 nibbles.
Registration without RotateQuant and accumulator-overflow decline coverage.
Three real Krea2 captured-activation cases: exact equality against eager using the revised production implementation.
Full Krea2 1024x1024 workflow (seed 20260907, 8 Euler/simple steps, CFG 1): the revised production implementation produced pixels identical to the eager baseline (MAE=0, SSIM=1). One diagnostic run took 346.34 s; this includes startup/loading effects and is not a repeated hot benchmark.
ruff check, ruff format --check, and git diff --check passed.
Relevant Ascend/backend/ConvRot/INT8/input-activation/constraint regression suite: 184 passed, 147 skipped, 1 warning. Skipped tests require other backends/configurations; the warning is the existing NPU internal-format allocation warning.

The ComfyUI validation checkout retains the existing local integration shim routing the custom node's mixed INT8 layers through Comfy Kitchen. It is identical between comparisons and is not part of this PR. No multi-seed quality, peak-memory, or general performance guarantee is implied by the fixed-seed test.

_{🔄 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/161 **Author:** [@Big2Wheel](https://github.com/Big2Wheel) **Created:** 9/7/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `codex/ascend-convrot-w4a4` --- ### 📝 Commits (10+) - [`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 - [`9ad6f5e`](https://github.com/Comfy-Org/comfy-kitchen/commit/9ad6f5e6dc416e8964ed45bbc2801125de985920) feat(ascend): add ConvRot W4A4 linear backend - [`aaad462`](https://github.com/Comfy-Org/comfy-kitchen/commit/aaad462eb856f6252d21058a15ed52ae79f88c76) fix(ascend): guard RotateQuant feature limit - [`8232c93`](https://github.com/Comfy-Org/comfy-kitchen/commit/8232c9357f47231ac7f097564ed93464797968cc) test(ascend): cover fallback without RotateQuant ### 📊 Changes **7 files changed** (+1368 additions, -45 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+76 -38) 📝 `comfy_kitchen/__init__.py` (+1 -0) ➕ `comfy_kitchen/backends/ascend/__init__.py` (+595 -0) 📝 `comfy_kitchen/registry.py` (+1 -1) 📝 `tests/conftest.py` (+2 -2) ➕ `tests/test_ascend_backend.py` (+668 -0) 📝 `tests/test_backends.py` (+25 -4) </details> ### 📄 Description ## What this PR does / why we need it Adds an optional Ascend implementation of `convrot_w4a4_linear`, prioritizing the eager reference's numerical behavior. ### Correctness revision The previous implementation silently converted FP32 activations to BF16 before rotation/quantization, and narrowed all matmul results through FP16. An audit demonstrated changed INT4 activation codes, fixed-seed image divergence, and finite FP32/BF16 inputs producing non-finite outputs where eager remained finite. This revision: - preserves the input dtype during rotation and reuses the eager signed-INT4 quantizer (including the [-7,7] emission range); - unpacks the original A4/W4 integer codes to INT8 and calls `npu_quant_matmul` with a unit FLOAT scale, no per-token scale, and INT32 output; - applies the cast, activation scale, weight scale, and bias in eager's order, with no compulsory FP16 intermediate for FP32/BF16; - rejects feature widths that could overflow the INT32 accumulator; - no longer requires RotateQuant for this capability; existing INT8/RotateQuant APIs and their width guards are unchanged; - retains fallback through the registry for unsupported call contracts. This is **not a zero-copy packed-INT4 matmul implementation**. It trades unpacking and reference-preprocessing cost for correctness. FP16 retains eager's own intermediate-range limitations; the implementation does not promise greater precision than the reference. ### Correction of earlier claims Both eager and Ascend paths quantize activations to A4. The earlier explanation that the image difference was expected simply because the optimized path quantizes activations was incorrect. The previously reported 134.00 s / 2.73x result belongs to the superseded, numerically different implementation. It is withdrawn as a performance claim for the current revision. New diagnostic timings are not a repeated hot A/B benchmark and no speedup is claimed here. ### Relationship to the other Ascend PRs - #151 supplies the base optional Ascend backend/registration. - #153 supplies the quantized-matmul backend infrastructure in the same module. This implementation does not call the public `int8_linear` API or change its semantics. - This branch remains stacked on #155 for review/history, but W4A4 no longer calls its fused RotateQuant helper. The Git stack dependency is distinct from the runtime requirement. - #154 (RoPE/RMS-RoPE) is orthogonal and is not included in this stack. After the predecessor stack lands, this PR can be rebased on `main` to leave the W4A4-specific changes. The branch also merges main through `21003fa`, retaining both the upstream allocation-context import and optional Ascend registration to resolve their import conflict. The same regression suite passed again after this merge. ## Does this PR introduce any user-facing change? Yes. Eligible Ascend ConvRot W4A4 calls use the correctness-first integer-accumulation implementation. Weights remain INT4-quantized on disk/in the caller, but the backend materializes INT8 values for matmul. No new mandatory dependency or opt-in environment variable is introduced. ## How was this patch tested? Hardware/runtime: one Ascend 910B3 (physical card 4), PyTorch 2.10.0+cpu with torch_npu 2.10.0.post4, CANN 9.1.0. - Zero-tolerance comparisons against eager for FP32, BF16 and FP16, including K=6144 and K=16384. - Large finite FP32/BF16 outputs beyond the FP16 range, zero/tiny/non-finite rows, bias, batched strided inputs, and signed/unsigned packed storage. - Operator-spy verification of original quantized codes, INT8 inputs, INT32 output, unit scale, no per-token scale, and no weight mutation. - Independent CPU INT64 oracle for the native integer accumulator, including negative INT4 nibbles. - Registration without RotateQuant and accumulator-overflow decline coverage. - Three real Krea2 captured-activation cases: exact equality against eager using the revised production implementation. - Full Krea2 1024x1024 workflow (seed 20260907, 8 Euler/simple steps, CFG 1): the revised production implementation produced pixels identical to the eager baseline (MAE=0, SSIM=1). One diagnostic run took 346.34 s; this includes startup/loading effects and is not a repeated hot benchmark. - `ruff check`, `ruff format --check`, and `git diff --check` passed. - Relevant Ascend/backend/ConvRot/INT8/input-activation/constraint regression suite: **184 passed, 147 skipped, 1 warning**. Skipped tests require other backends/configurations; the warning is the existing NPU internal-format allocation warning. The ComfyUI validation checkout retains the existing local integration shim routing the custom node's mixed INT8 layers through Comfy Kitchen. It is identical between comparisons and is not part of this PR. No multi-seed quality, peak-memory, or general performance guarantee is implied by the fixed-seed test. --- <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:22 +00:00