[PR #212] [Perf][Ascend] Reuse unpacked W4A4 weights within an explicit scope #224

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/212
Author: @Big2Wheel
Created: 9/29/2026
Status: 🔄 Open

Base: main ← Head: ascend-w4a4-weight-reuse

📝 Commits (2)

c608666 perf(ascend): add scoped content-validated W4A4 weight reuse
0f2fbf9 fix(ascend): preserve strided W4A4 cache identity and accept uint8

📊 Changes

9 files changed (+498 additions, -2 deletions)

View changed files

📝 comfy_kitchen/backends/ascend/__init__.py (+12 -2)
➕ comfy_kitchen/backends/ascend/weight_cache.py (+129 -0)
📝 comfy_kitchen/backends/cuda/__init__.py (+2 -0)
📝 comfy_kitchen/backends/eager/convrot_w4a4.py (+2 -0)
📝 comfy_kitchen/backends/hip/__init__.py (+2 -0)
📝 comfy_kitchen/tensor/convrot_w4a4.py (+4 -0)
➕ samples/ascend_w4a4_weight_cache.md (+44 -0)
➕ samples/benchmark_ascend_w4a4_weight_cache.py (+81 -0)
➕ tests/test_ascend_weight_cache.py (+222 -0)

📄 Description

Summary

Reuse unpacked Ascend W4A4 weights within an explicit caller-owned scope. Repeated denoising steps currently unpack the same packed weights before each integer matmul. This PR only caches that weight preparation; it does not cache activations, scales, outputs, or change quantization math.

Add ascend.W4A4WeightCache(max_bytes) and a keyword-only weight_cache argument to the public ConvRot W4A4 Linear dispatcher.
Keep default behavior unchanged and preserve backend selection. Other built-in backends accept and ignore the optional optimization.
Validate cache hits with packed-content equality, including inference tensors, .data/alias writes and storage replacement. This synchronization is included in the measurements.
Bound storage, weakly reference source weights, clear on normal/exceptional scope exit, and bypass reuse for cross-thread/stream calls and graph capture.
Include tests, a synchronized alternating Linear microbenchmark, and scope/memory documentation.

This PR is based directly on Kitchen main 888b13e and does not include or require #201. #201 removes activation pack/unpack; this PR reuses weight unpacking. They are complementary and should not be credited as the same speedup.

Companion ComfyUI draft: https://github.com/Comfy-Org/ComfyUI/pull/16640, which owns one cache per sampling operation. Kitchen alone does not automatically enable reuse in existing workflows. Merge/release Kitchen first, then update the ComfyUI dependency pin before merging that integration.

Validation

Single Ascend 910B3; Python 3.12.13, PyTorch 2.10.0 + torch-npu 2.10.0.post4, CANN 9.1.0 / driver 25.2.2.

python -m pytest -q tests/test_ascend_weight_cache.py tests/test_ascend_backend.py tests/test_convrot_w4a4.py tests/test_backends.py

146 passed, 41 skipped (CUDA-dependent cases); NPU cache/mutation/stream/Linear tests executed. ruff check . and git diff --check pass. The initial review found uint8 rejection and lost cache identity for strided weights; five regression cases failed before the fix and pass afterward.

Workflow validation uses Krea2 mixed INT4: krea2_turbo_int4_mixed.safetensors, qwen3vl_4b_bf16.safetensors, and qwen_image_vae.safetensors; OTUNetLoaderW4A4 from ComfyUI-INT4-Fast, no on-the-fly quantization or LoRA. 1024×1024, 8 steps, CFG 1, Euler/simple, seed 20260907, intermediate node caching disabled. Both A/B arms use the same existing NPU INT8 compatibility environment; this PR does not remove that separate dependency.

Latest-main integration smoke test (ComfyUI 67f11cc + companion changes; Kitchen 888b13e + this PR, without #201):

Cache limit	First request	Second request (hot)
Disabled	347.839 s	345.990 s
16 GiB	328.058 s	327.415 s

The single hot comparison is 5.37% lower, not a statistically established speedup (n=1 hot request per arm, separate service processes). This final rebase check was shortened from the longer original protocol because each request takes over five minutes without #201. Outputs have identical decoded RGB bytes across all four requests. Each cached request has 128 misses, 896 hits and 14022 MiB logical storage, including construction in the workflow time. Results were measured on c608666; subsequent review fixes preserve the tested contiguous-int8 path and add uint8/strided regression coverage.

The standalone-main gain is modest relative to the memory cost. The larger result below is conditional on the earlier #201-based environment, not a standalone claim or a directly comparable combined speedup across source revisions.

Earlier integration evidence (contains #201; not the standalone-main baseline)

On ComfyUI 45b3121 and Kitchen #201 a40a0db plus this cache, one warmup then three hot runs per arm gave 36.367 s disabled vs 16.461 s enabled (54.74% lower). This was grouped integration confirmation, not alternating statistical benchmarking. Decoded RGB outputs were byte-identical. Every cached request rebuilt its cache: 128 misses, 896 hits, 13.693 GiB logical storage. Model switching Krea → Z-Image BF16 → Krea, interruption with a populated cache followed by a successful next request, and a 1 MiB budget with zero admitted weights all completed correctly.

Memory and support limits

This is an explicit memory-for-time tradeoff, not a low-memory optimization. Storage includes the packed validation snapshot plus decoded INT8 data (~3× packed bytes). The caller must reserve headroom; budget overflow uses uncached unpacking without live-entry eviction, but allocator OOM is not caught/retried. Concurrent source mutation is unsupported. No cross-request/global cache, checkpoint changes, or automatic enablement.

_{🔄 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/212 **Author:** [@Big2Wheel](https://github.com/Big2Wheel) **Created:** 9/29/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `ascend-w4a4-weight-reuse` --- ### 📝 Commits (2) - [`c608666`](https://github.com/Comfy-Org/comfy-kitchen/commit/c608666eea111c7d174008f70c7f8eb8a733d8ba) perf(ascend): add scoped content-validated W4A4 weight reuse - [`0f2fbf9`](https://github.com/Comfy-Org/comfy-kitchen/commit/0f2fbf9ef2e5ac88bd1690f56df6a3b637b9b248) fix(ascend): preserve strided W4A4 cache identity and accept uint8 ### 📊 Changes **9 files changed** (+498 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/ascend/__init__.py` (+12 -2) ➕ `comfy_kitchen/backends/ascend/weight_cache.py` (+129 -0) 📝 `comfy_kitchen/backends/cuda/__init__.py` (+2 -0) 📝 `comfy_kitchen/backends/eager/convrot_w4a4.py` (+2 -0) 📝 `comfy_kitchen/backends/hip/__init__.py` (+2 -0) 📝 `comfy_kitchen/tensor/convrot_w4a4.py` (+4 -0) ➕ `samples/ascend_w4a4_weight_cache.md` (+44 -0) ➕ `samples/benchmark_ascend_w4a4_weight_cache.py` (+81 -0) ➕ `tests/test_ascend_weight_cache.py` (+222 -0) </details> ### 📄 Description ## Summary Reuse unpacked Ascend W4A4 weights within an explicit caller-owned scope. Repeated denoising steps currently unpack the same packed weights before each integer matmul. This PR only caches that weight preparation; it does not cache activations, scales, outputs, or change quantization math. - Add `ascend.W4A4WeightCache(max_bytes)` and a keyword-only `weight_cache` argument to the public ConvRot W4A4 Linear dispatcher. - Keep default behavior unchanged and preserve backend selection. Other built-in backends accept and ignore the optional optimization. - Validate cache hits with packed-content equality, including inference tensors, `.data`/alias writes and storage replacement. This synchronization is included in the measurements. - Bound storage, weakly reference source weights, clear on normal/exceptional scope exit, and bypass reuse for cross-thread/stream calls and graph capture. - Include tests, a synchronized alternating Linear microbenchmark, and scope/memory documentation. ## Related work This PR is based directly on Kitchen main `888b13e` and **does not include or require #201**. #201 removes **activation** pack/unpack; this PR reuses **weight** unpacking. They are complementary and should not be credited as the same speedup. Companion ComfyUI draft: https://github.com/Comfy-Org/ComfyUI/pull/16640, which owns one cache per sampling operation. Kitchen alone does not automatically enable reuse in existing workflows. Merge/release Kitchen first, then update the ComfyUI dependency pin before merging that integration. ## Validation Single Ascend 910B3; Python 3.12.13, PyTorch 2.10.0 + torch-npu 2.10.0.post4, CANN 9.1.0 / driver 25.2.2. ```bash python -m pytest -q tests/test_ascend_weight_cache.py tests/test_ascend_backend.py tests/test_convrot_w4a4.py tests/test_backends.py ``` **146 passed, 41 skipped** (CUDA-dependent cases); NPU cache/mutation/stream/Linear tests executed. `ruff check .` and `git diff --check` pass. The initial review found uint8 rejection and lost cache identity for strided weights; five regression cases failed before the fix and pass afterward. Workflow validation uses Krea2 mixed INT4: `krea2_turbo_int4_mixed.safetensors`, `qwen3vl_4b_bf16.safetensors`, and `qwen_image_vae.safetensors`; `OTUNetLoaderW4A4` from ComfyUI-INT4-Fast, no on-the-fly quantization or LoRA. 1024×1024, 8 steps, CFG 1, Euler/simple, seed 20260907, intermediate node caching disabled. Both A/B arms use the same existing NPU INT8 compatibility environment; this PR does not remove that separate dependency. Latest-main integration smoke test (ComfyUI `67f11cc` + companion changes; Kitchen `888b13e` + this PR, **without #201**): | Cache limit | First request | Second request (hot) | | --- | ---: | ---: | | Disabled | 347.839 s | 345.990 s | | 16 GiB | 328.058 s | 327.415 s | The single hot comparison is **5.37% lower**, not a statistically established speedup (n=1 hot request per arm, separate service processes). This final rebase check was shortened from the longer original protocol because each request takes over five minutes without #201. Outputs have identical decoded RGB bytes across all four requests. Each cached request has 128 misses, 896 hits and 14022 MiB logical storage, including construction in the workflow time. Results were measured on `c608666`; subsequent review fixes preserve the tested contiguous-int8 path and add uint8/strided regression coverage. **The standalone-main gain is modest relative to the memory cost.** The larger result below is conditional on the earlier #201-based environment, not a standalone claim or a directly comparable combined speedup across source revisions. <details> <summary>Earlier integration evidence (contains #201; not the standalone-main baseline)</summary> On ComfyUI `45b3121` and Kitchen #201 `a40a0db` plus this cache, one warmup then three hot runs per arm gave 36.367 s disabled vs 16.461 s enabled (54.74% lower). This was grouped integration confirmation, not alternating statistical benchmarking. Decoded RGB outputs were byte-identical. Every cached request rebuilt its cache: 128 misses, 896 hits, 13.693 GiB logical storage. Model switching Krea → Z-Image BF16 → Krea, interruption with a populated cache followed by a successful next request, and a 1 MiB budget with zero admitted weights all completed correctly. </details> ## Memory and support limits This is an explicit memory-for-time tradeoff, **not a low-memory optimization**. Storage includes the packed validation snapshot plus decoded INT8 data (~3× packed bytes). The caller must reserve headroom; budget overflow uses uncached unpacking without live-entry eviction, but allocator OOM is not caught/retried. Concurrent source mutation is unsupported. No cross-request/global cache, checkpoint changes, or automatic enablement. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

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

bsakaguchi referenced this issue

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

[PR #227] Tune unmasked SM120 D64 attention with smaller query tiles #236

bsakaguchi referenced this issue

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

[PR #224] Fix nonpositive scale ordering in CUDA INT8 attention #239

bsakaguchi referenced this issue

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