[PR #201] [MERGED] [Perf][Ascend] Avoid activation pack/unpack in ConvRot W4A4 #218

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/201
Author: @Big2Wheel
Created: 9/25/2026
Status: ✅ Merged
Merged: 10/2/2026
Merged by: @comfyanonymous

Base: main ← Head: perf/ascend-w4a4-unpacked-activations

📝 Commits (4)

84f29d0 fix(ascend): sort backend imports to unblock CI
a40a0db perf(ascend): avoid W4A4 activation pack and unpack
e72ec8b Fix import order.
59ed9c0 Fix

📊 Changes

6 files changed (+227 additions, -10 deletions)

View changed files

📝 README.md (+3 -2)
📝 comfy_kitchen/backends/ascend/__init__.py (+9 -6)
📝 comfy_kitchen/backends/eager/convrot_w4a4.py (+11 -2)
➕ samples/benchmark_ascend_w4a4_unpacked.py (+92 -0)
📝 tests/test_ascend_backend.py (+15 -0)
➕ tests/test_convrot_unpacked_quantization.py (+97 -0)

📄 Description

Summary

Avoid packing signed INT4 activation codes into nibbles only to immediately unpack them before Ascend ConvRot W4A4 integer matmul. Keep the same codes in INT8 storage instead.

Extract a private unpacked quantization helper; preserve the public packed quantizer, its signature, scales, and stochastic-rounding behavior.
Keep activation rotation precision, rounding, weight unpacking, INT32 accumulation, and cast/scale/bias ordering unchanged.
Add exact-equivalence tests, a regression assertion that only weights are unpacked, and a standalone alternating-arm NPU benchmark.
No packed A4W4 GEMM substitution, precision reduction, weight cache, or new dependency.

Relationship to #198 and #199

#198 (merged) is the functional baseline. This optimizes its correctness-preserving Ascend W4A4 implementation, including the numerical fixes originally developed in #161. Both benchmark arms use the same native INT32 npu_quant_matmul; the baseline is not eager/CPU GEMM or the earlier withdrawn lower-precision implementation.
#199 (currently open) is an independent maintenance follow-up: RoPE capability-gated tests, support documentation, and backend import ordering. It does not contain this runtime optimization, and this PR does not replace it.
This branch targets main directly, not #199. To avoid the existing main-branch Ruff I001 failure, its first commit cherry-picks only #199's import-order fix (28b7f850cbc4edc0e3e0605e026252c5bd3a53c1, with provenance). No other #199 changes are included. The second commit is the performance change. Prefer merging #199 first; this shared import-only prerequisite can then be dropped when rebasing. The runtime optimization has no dependency on #199's test/documentation changes.

Correctness and tests

NPU regression selection: 345 passed, 172 skipped (other hardware/configurations); two existing format/reference warnings.
CPU-only optional-import simulation with torch-npu unavailable: 113 passed, 79 skipped.
Exact packed/unpacked code and scale comparisons cover FP32/FP16/BF16, contiguous/strided/transposed inputs, rounding seeds, zeros, tiny/max-finite values, NaN/Inf, rounding boundaries, empty rows, and preservation of the public odd-width rejection.
Three captured real Krea activation/weight samples with K=6144/16384: new = old = eager with rtol=0, atol=0.
Full-repository ruff check . and git diff --check pass.

NPU regression command (run in the source checkout):

ASCEND_RT_VISIBLE_DEVICES=4 OMP_NUM_THREADS=4 PYTHONPATH=. python -m pytest -q \
  tests/test_convrot_unpacked_quantization.py tests/test_ascend_backend.py \
  tests/test_ascend_rope.py tests/test_convrot_w4a4.py tests/test_int8.py \
  tests/test_int8_input_act.py tests/test_constraints.py tests/test_backends.py

Operator performance and profiling

Environment: one Ascend 910B3 (physical device 4), torch 2.10.0+cpu, torch-npu 2.10.0.post4, CANN 9.1.0, Python 3.12.13, vLLM-Omni v0.28.0 container. Runs were isolated from other test jobs on this card.

FP32 full Linear timings include rotation, activation preparation, weight unpacking, matmul, and postprocessing. Random full-shape inputs; 10 warmups and 30 measurements per arm, alternating order, with NPU synchronization at timing boundaries. Real captured inputs were checked separately for exact correctness.

M / N / K	Baseline median (ms)	Unpacked median (ms)	Peak allocation increment, baseline/new (MiB)
4135 / 6144 / 6144	235.283	12.740	697.163 / 697.163
4135 / 16384 / 6144	255.638	33.360	1657.163 / 1657.163
4135 / 6144 / 16384	580.379	35.111	1859.067 / 1859.067

Quantization/preparation alone (10 warmups, 50 measurements per arm) was 224.920 → 0.957 ms for [4135,6144], and 549.891 → 3.189 ms for [4135,16384].

Separate profiling confirms removal of activation aten::__lshift__/aten::bitwise_or; weight unpacking and one npu::npu_quant_matmul remain. A scheduled full-shape diagnostic profile for [4135,6144] measured the pack region at 218.555 ms, dominated by the aclnnLeftShifts_LeftShift_LeftShift event at 210.135 ms. These single-profile numbers explain the hotspot; they are not the steady-state benchmark. Device events alone are not evidence of host CPU fallback or a specific accelerator core type.

Reproduce the full Linear comparison in a standalone process:

ASCEND_RT_VISIBLE_DEVICES=4 OMP_NUM_THREADS=4 PYTHONPATH=. python \
  samples/benchmark_ascend_w4a4_unpacked.py \
  --m 4135 --n 6144 --k 6144 --dtype float32 --warmup 10 --iterations 30

The script checks exact output equality and prints raw samples and medians. Its legacy arm restores only activation pack/unpack through a temporary private-helper replacement.

Real Krea workflow A/B

ComfyUI commit: 45b3121ccff27bc6020985738ad999b96b4b9a02.
Model: krea2_turbo_int4_mixed.safetensors; text encoder: qwen3vl_4b_bf16.safetensors; VAE: qwen_image_vae.safetensors.
1024×1024, batch 1, 8 steps, CFG 1, Euler/simple; seed 20260907, additional correctness seed 20260926.
Existing ComfyUI-INT4-Fast custom node and its local mixed-INT8 compatibility shim were held fixed. This PR does not remove that shim dependency or claim this workflow is fully upstream-only.
--cache-none --disable-pinned-memory, OMP_NUM_THREADS=4; same service, same instrumentation in both arms, only the activation quantization helper changes.
One warmup per arm, then measured order A/B/B/A/A/B (three runs each). All runs executed 1024 W4A4 calls: 832 with activation shape [4135,6144] and 192 with [4135,16384]. No node-cache shortcut.
Workflow time is server history execution_start → execution_success, including loading, encoding, sampling, decoding, and saving, but excluding service startup/client polling. Denoising is synchronized around KSAMPLER.sample, not pure kernel time. Runtime/filesystem are warm; model-load calls still occur, so this is not a fully resident-model claim.

Metric	Baseline runs (s)	Unpacked runs (s)	Baseline/new median (s)
Full workflow	330.316, 330.694, 329.570	36.037, 36.116, 35.909	330.316 / 36.037
Denoising	325.605, 325.522, 324.685	31.005, 30.993, 31.014	325.522 / 31.005

Workflow mean ± sample standard deviation: 330.193 ± 0.572 s → 36.021 ± 0.104 s. Median latency reduction: 89.09% (9.17× speedup) in this specific environment/workflow. This is not a universal model/hardware speedup claim; only three formal runs per arm were measured.

Warmups: 330.760 / 35.997 s (excluded from statistics). Additional seed: 330.784 / 35.982 s (also excluded). All 10 requests succeeded; decoded output images were pixel-identical across arms for each seed. No CPU-fallback warnings or execution errors were observed in the workflow log; that is not an exhaustive per-operator device audit.

Peak allocated memory during formal runs: baseline 11247.688 MiB, new 11247.346 MiB. No significant peak-memory reduction is claimed. Z-Image/Wan results from #199 are separate validation of #198 and are not presented as measurements of this change.

_{🔄 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/201 **Author:** [@Big2Wheel](https://github.com/Big2Wheel) **Created:** 9/25/2026 **Status:** ✅ Merged **Merged:** 10/2/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `perf/ascend-w4a4-unpacked-activations` --- ### 📝 Commits (4) - [`84f29d0`](https://github.com/Comfy-Org/comfy-kitchen/commit/84f29d0ee58fcfe84263a74e607b3da330c65e49) fix(ascend): sort backend imports to unblock CI - [`a40a0db`](https://github.com/Comfy-Org/comfy-kitchen/commit/a40a0dbadba1f8b91f46ed035953b631f73837c3) perf(ascend): avoid W4A4 activation pack and unpack - [`e72ec8b`](https://github.com/Comfy-Org/comfy-kitchen/commit/e72ec8b98cb0f84dc12e6ea9a5a64793eaf3c243) Fix import order. - [`59ed9c0`](https://github.com/Comfy-Org/comfy-kitchen/commit/59ed9c01443d9247a8407a1cfa56fc8107ca5b99) Fix ### 📊 Changes **6 files changed** (+227 additions, -10 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+3 -2) 📝 `comfy_kitchen/backends/ascend/__init__.py` (+9 -6) 📝 `comfy_kitchen/backends/eager/convrot_w4a4.py` (+11 -2) ➕ `samples/benchmark_ascend_w4a4_unpacked.py` (+92 -0) 📝 `tests/test_ascend_backend.py` (+15 -0) ➕ `tests/test_convrot_unpacked_quantization.py` (+97 -0) </details> ### 📄 Description ## Summary Avoid packing signed INT4 activation codes into nibbles only to immediately unpack them before Ascend ConvRot W4A4 integer matmul. Keep the same codes in INT8 storage instead. - Extract a private unpacked quantization helper; preserve the public packed quantizer, its signature, scales, and stochastic-rounding behavior. - Keep activation rotation precision, rounding, weight unpacking, INT32 accumulation, and cast/scale/bias ordering unchanged. - Add exact-equivalence tests, a regression assertion that only weights are unpacked, and a standalone alternating-arm NPU benchmark. - No packed A4W4 GEMM substitution, precision reduction, weight cache, or new dependency. ## Relationship to #198 and #199 - **#198 (merged) is the functional baseline.** This optimizes its correctness-preserving Ascend W4A4 implementation, including the numerical fixes originally developed in #161. Both benchmark arms use the same native INT32 `npu_quant_matmul`; the baseline is not eager/CPU GEMM or the earlier withdrawn lower-precision implementation. - **#199 (currently open) is an independent maintenance follow-up:** RoPE capability-gated tests, support documentation, and backend import ordering. It does not contain this runtime optimization, and this PR does not replace it. - This branch targets `main` directly, not #199. To avoid the existing main-branch Ruff I001 failure, its first commit cherry-picks only #199's import-order fix (`28b7f850cbc4edc0e3e0605e026252c5bd3a53c1`, with provenance). No other #199 changes are included. The second commit is the performance change. Prefer merging #199 first; this shared import-only prerequisite can then be dropped when rebasing. The runtime optimization has no dependency on #199's test/documentation changes. ## Correctness and tests - NPU regression selection: **345 passed, 172 skipped** (other hardware/configurations); two existing format/reference warnings. - CPU-only optional-import simulation with torch-npu unavailable: **113 passed, 79 skipped**. - Exact packed/unpacked code and scale comparisons cover FP32/FP16/BF16, contiguous/strided/transposed inputs, rounding seeds, zeros, tiny/max-finite values, NaN/Inf, rounding boundaries, empty rows, and preservation of the public odd-width rejection. - Three captured real Krea activation/weight samples with K=6144/16384: new = old = eager with `rtol=0, atol=0`. - Full-repository `ruff check .` and `git diff --check` pass. NPU regression command (run in the source checkout): ```bash ASCEND_RT_VISIBLE_DEVICES=4 OMP_NUM_THREADS=4 PYTHONPATH=. python -m pytest -q \ tests/test_convrot_unpacked_quantization.py tests/test_ascend_backend.py \ tests/test_ascend_rope.py tests/test_convrot_w4a4.py tests/test_int8.py \ tests/test_int8_input_act.py tests/test_constraints.py tests/test_backends.py ``` ## Operator performance and profiling Environment: one Ascend 910B3 (physical device 4), torch 2.10.0+cpu, torch-npu 2.10.0.post4, CANN 9.1.0, Python 3.12.13, vLLM-Omni v0.28.0 container. Runs were isolated from other test jobs on this card. FP32 full Linear timings include rotation, activation preparation, weight unpacking, matmul, and postprocessing. Random full-shape inputs; 10 warmups and 30 measurements per arm, alternating order, with NPU synchronization at timing boundaries. Real captured inputs were checked separately for exact correctness. | M / N / K | Baseline median (ms) | Unpacked median (ms) | Peak allocation increment, baseline/new (MiB) | | --- | ---: | ---: | ---: | | 4135 / 6144 / 6144 | 235.283 | 12.740 | 697.163 / 697.163 | | 4135 / 16384 / 6144 | 255.638 | 33.360 | 1657.163 / 1657.163 | | 4135 / 6144 / 16384 | 580.379 | 35.111 | 1859.067 / 1859.067 | Quantization/preparation alone (10 warmups, 50 measurements per arm) was 224.920 → 0.957 ms for [4135,6144], and 549.891 → 3.189 ms for [4135,16384]. Separate profiling confirms removal of activation `aten::__lshift__`/`aten::bitwise_or`; weight unpacking and one `npu::npu_quant_matmul` remain. A scheduled full-shape diagnostic profile for [4135,6144] measured the pack region at 218.555 ms, dominated by the `aclnnLeftShifts_LeftShift_LeftShift` event at 210.135 ms. These single-profile numbers explain the hotspot; they are not the steady-state benchmark. Device events alone are not evidence of host CPU fallback or a specific accelerator core type. Reproduce the full Linear comparison in a standalone process: ```bash ASCEND_RT_VISIBLE_DEVICES=4 OMP_NUM_THREADS=4 PYTHONPATH=. python \ samples/benchmark_ascend_w4a4_unpacked.py \ --m 4135 --n 6144 --k 6144 --dtype float32 --warmup 10 --iterations 30 ``` The script checks exact output equality and prints raw samples and medians. Its legacy arm restores only activation pack/unpack through a temporary private-helper replacement. ## Real Krea workflow A/B - ComfyUI commit: `45b3121ccff27bc6020985738ad999b96b4b9a02`. - Model: `krea2_turbo_int4_mixed.safetensors`; text encoder: `qwen3vl_4b_bf16.safetensors`; VAE: `qwen_image_vae.safetensors`. - 1024×1024, batch 1, 8 steps, CFG 1, Euler/simple; seed 20260907, additional correctness seed 20260926. - Existing ComfyUI-INT4-Fast custom node and its local mixed-INT8 compatibility shim were held fixed. **This PR does not remove that shim dependency or claim this workflow is fully upstream-only.** - `--cache-none --disable-pinned-memory`, `OMP_NUM_THREADS=4`; same service, same instrumentation in both arms, only the activation quantization helper changes. - One warmup per arm, then measured order A/B/B/A/A/B (three runs each). All runs executed 1024 W4A4 calls: 832 with activation shape [4135,6144] and 192 with [4135,16384]. No node-cache shortcut. - Workflow time is server history execution_start → execution_success, including loading, encoding, sampling, decoding, and saving, but excluding service startup/client polling. Denoising is synchronized around KSAMPLER.sample, not pure kernel time. Runtime/filesystem are warm; model-load calls still occur, so this is not a fully resident-model claim. | Metric | Baseline runs (s) | Unpacked runs (s) | Baseline/new median (s) | | --- | --- | --- | --- | | Full workflow | 330.316, 330.694, 329.570 | 36.037, 36.116, 35.909 | 330.316 / 36.037 | | Denoising | 325.605, 325.522, 324.685 | 31.005, 30.993, 31.014 | 325.522 / 31.005 | Workflow mean ± sample standard deviation: 330.193 ± 0.572 s → 36.021 ± 0.104 s. Median latency reduction: **89.09% (9.17× speedup) in this specific environment/workflow**. This is not a universal model/hardware speedup claim; only three formal runs per arm were measured. Warmups: 330.760 / 35.997 s (excluded from statistics). Additional seed: 330.784 / 35.982 s (also excluded). All 10 requests succeeded; decoded output images were pixel-identical across arms for each seed. No CPU-fallback warnings or execution errors were observed in the workflow log; that is not an exhaustive per-operator device audit. Peak allocated memory during formal runs: baseline 11247.688 MiB, new 11247.346 MiB. **No significant peak-memory reduction is claimed.** Z-Image/Wan results from #199 are separate validation of #198 and are not presented as measurements of this change. --- <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:41 +00:00