[PR #198] [MERGED] [Feat][Ascend] Consolidate INT8, RotateQuant, RoPE and ConvRot W4A4 backends #215

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: codex/ascend-consolidated

📝 Commits (6)

3bde6fd feat(ascend): add INT8 linear backend
b8b5692 feat(ascend): add guarded fused RotateQuant support
a2224f3 feat(ascend): add RoPE and RMS-RoPE with independent capabilities
e42e10b feat(ascend): add numerically faithful ConvRot W4A4 linear
cb04fec fix(ascend): match current linear API and test capability isolation
8cb9fda docs(ascend): explain consolidated capability requirements

📊 Changes

6 files changed (+1957 additions, -47 deletions)

View changed files

📝 README.md (+48 -12)
📝 comfy_kitchen/backends/ascend/__init__.py (+643 -26)
➕ comfy_kitchen/backends/ascend/rope.py (+365 -0)
📝 tests/test_ascend_backend.py (+601 -1)
➕ tests/test_ascend_rope.py (+292 -0)
📝 tests/test_backends.py (+8 -8)

📄 Description

Summary and relationship to the previous PRs

Following the suggestion to consolidate the remaining Ascend work after #151 was merged, this PR starts directly from main (d86e8d7) and replaces the conflicting feature stack:

Previous PR	Included functionality
#153	Dynamically quantized INT8 linear using `npu_quant_matmul`
#154	Interleaved/split-half RoPE and RMS-RoPE, single/paired and in-place APIs
#155	Optional fused RotateQuant for supported INT8 ConvRot calls
#161	Correctness-first ConvRot W4A4 integer-accumulation backend

The base backend from #151 is already in main and is not reintroduced. Feature implementations remain separate signed-off commits, followed by current-main compatibility fixes and documentation. W4A4 does not call the public INT8 linear API or require RotateQuant; RoPE is independent of the quantization features.

Integration changes

Resolve the overlapping backend, capability table and test changes against main.
Preserve independent RoPE/RMS-RoPE capability detection, including operator-schema checks for rotary_mode and div_mode.
Preserve optional RotateQuant fallback, zero-scale handling and the K > 16000 fused-kernel guard.
Adapt INT8 linear to main's new input_act_weight, input_act_eps, residual and residual_scale arguments. Input RMSNorm and residual modulation use the shared device-side helpers, not a new fused GEMM implementation.
Test that missing quantization/matmul/RotateQuant capabilities do not disable unrelated functionality.
Preserve packed/strided RoPE views, in-place alias checks, partial RMS-RoPE and registry fallback for unsupported contracts.

W4A4 correctness and performance scope

This includes the corrected #161 implementation, not the earlier numerically different packed-kernel path:

Preserve FP32/FP16/BF16 activation dtype during rotation and use the reference signed-INT4 quantizer.
Unpack the same A4/W4 codes to INT8 and accumulate with npu_quant_matmul into INT32, using a unit scale and no per-token scale.
Cast and apply activation scale, weight scale and bias in eager's order; guard possible INT32 accumulator overflow.

There is no compulsory FP32-to-BF16 input conversion or FP16 intermediate for FP32/BF16 outputs. FP16 retains the reference's own range limitations. This is not zero-copy packed-INT4 GEMM: unpacking adds memory/runtime cost.

The earlier 134 s / 2.73x result from the superseded W4A4 implementation is not a performance claim for this PR. The integration workflow runs below are smoke tests, not repeated hot A/B measurements, and establish no new speedup or multi-seed quality guarantee.

User-facing changes

Eligible NPU calls can use the above Ascend implementations automatically through the existing registry. No mandatory MindIE-SD dependency, environment variable or public API change is introduced by this PR. Non-NPU import remains optional; unsupported contracts continue through the registry's existing fallback behavior.

Validation

One Ascend 910B3, physical card 4; PyTorch 2.10.0+cpu, torch_npu 2.10.0.post4, CANN 9.1.0.

Relevant NPU/backend/INT8/ConvRot/RoPE/constraint regression: 277 passed, 172 skipped (two framework warnings). Skips cover other hardware/backend configurations.
CPU optional-import and backend/constraint/RoPE/RMS-RoPE regression with TORCH_DEVICE_BACKEND_AUTOLOAD=0 and sys.modules["torch_npu"] = None: 361 passed, 598 skipped, one warning. This explicitly simulates torch-npu being unavailable without uninstalling it from the shared validation container.
Replayed three saved real Krea activation/weight samples (from layers with K=6144 and K=16384): the integrated W4A4 implementation matches eager with rtol=0, atol=0 in every case.
Fresh single-card ComfyUI smoke tests with the integrated checkout and node cache disabled: Z-Image INT8 and BF16 (1920x1088, 8 steps), Wan2.2 14B (640x640, 81 frames, 4 steps), and Krea2 mixed INT4 (1024x1024, seed 20260907, 8 steps) all completed successfully. Dispatch tracing confirmed Ascend INT8 linear/RMS-RoPE, Wan RoPE and Krea W4A4. Unsupported Krea RoPE calls retained eager fallback.
The new Krea image is pixel-identical (MAE=0) to the saved eager baseline from the #161 correctness audit; this was not a new two-arm A/B benchmark. All 81 output Wan frames decoded successfully. Z-Image/Krea images were also visually inspected.
ruff check, ruff format --check on all changed Python files and git diff --check: passed.

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

The ComfyUI validation checkout retains the existing local custom-node shim routing mixed INT8 layers through Comfy Kitchen. That shim is not included in this PR. Workflow smoke results must not be interpreted as unmodified-ComfyUI support for every third-party custom node.

Completed workflows, exact weights and output specifications

The following are the exact checkpoint filenames selected in the saved successful API requests on 2026-09-25. All runs used one Ascend 910B3 (physical card 4), batch size 1, --cache-none --disable-pinned-memory, and OMP_NUM_THREADS=4.

Workflow	Main diffusion checkpoint(s)	Output (width × height)	Sampling	Observed request-to-completion time
Z-Image Turbo INT8 ConvRot	`z_image_turbo_int8_convrot.safetensors`	One image, 1920 × 1088	8 steps; CFG 1; res_multistep/simple	28.98 s
Z-Image Turbo BF16	`z_image_turbo_bf16.safetensors`	One image, 1920 × 1088	8 steps; CFG 1; res_multistep/simple	23.76 s
Wan2.2 T2V 14B FP16 + LightX2V	`wan2.2_t2v_high_noise_14B_fp16.safetensors` + `wan2.2_t2v_low_noise_14B_fp16.safetensors`	One video, 640 × 640; 81 frames at 16 FPS; 5.0625 s	4 steps total (2 high-noise + 2 low-noise); CFG 1; Euler/simple	137.11 s
Krea2 Turbo mixed INT4	`krea2_turbo_int4_mixed.safetensors`	One image, 1024 × 1024	8 steps; CFG 1; Euler/simple	356.64 s

Supporting weights used by these workflows:

Workflow	Text encoder	VAE	LoRA
Both Z-Image workflows	`qwen_3_4b.safetensors`	`ae.safetensors`	None
Wan2.2	`umt5_xxl_fp16.safetensors`	`wan_2.1_vae.safetensors`	`wan2.2_t2v_lightx2v_4steps_lora_v1.1_high_noise.safetensors` and `wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors`, each at strength 1.0 on its corresponding model
Krea2	`qwen3vl_4b_bf16.safetensors`	`qwen_image_vae.safetensors`	None

Seeds: both Z-Image runs used 493240645134828; Wan used 832596732254118 for initial noise (the low-noise stage did not add noise); Krea used 20260907. Krea used the existing ComfyUI-INT4-Fast custom node (OTUNetLoaderW4A4) and the local mixed-INT8 integration shim described above.

Timing definition and limitations:

These are single successful smoke-run wall times, measured immediately before submitting the API request until successful history completion was observed. They include workflow execution, any model loading/transfers performed by that request, VAE decoding, output saving/video encoding, and polling overhead (approximately one second polling interval); they exclude service startup.
The workflows ran sequentially in the same service in the table order. Node output caching was disabled, but model residency/allocator state can carry across requests. Other validation checks also ran during the session. These are not isolated, uniformly cold or repeated hot measurements.
Therefore the table is execution evidence, not an INT8-vs-BF16 comparison, an optimization speedup claim, or a performance guarantee.
Image dimensions were checked from the generated files. The Wan MP4 metadata reports 640 × 640, 81 frames, 16 FPS and 5.0625 seconds, and all 81 frames decoded successfully. Video playback duration is separate from the 137.11-second generation wall time.

_{🔄 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/198 **Author:** [@Big2Wheel](https://github.com/Big2Wheel) **Created:** 9/25/2026 **Status:** ✅ Merged **Merged:** 9/25/2026 **Merged by:** [@comfyanonymous](https://github.com/comfyanonymous) **Base:** `main` ← **Head:** `codex/ascend-consolidated` --- ### 📝 Commits (6) - [`3bde6fd`](https://github.com/Comfy-Org/comfy-kitchen/commit/3bde6fd1ae2c12017725e5b765dbb834efd1b939) feat(ascend): add INT8 linear backend - [`b8b5692`](https://github.com/Comfy-Org/comfy-kitchen/commit/b8b5692f6f251e8c98768a20a4e64a5150bdb933) feat(ascend): add guarded fused RotateQuant support - [`a2224f3`](https://github.com/Comfy-Org/comfy-kitchen/commit/a2224f37501883e0272a62de159335032254dfba) feat(ascend): add RoPE and RMS-RoPE with independent capabilities - [`e42e10b`](https://github.com/Comfy-Org/comfy-kitchen/commit/e42e10bc65d25a99153cc630bab9440588439b5f) feat(ascend): add numerically faithful ConvRot W4A4 linear - [`cb04fec`](https://github.com/Comfy-Org/comfy-kitchen/commit/cb04fecdfa73dc1fbedb8f01cf2d1cda808d9743) fix(ascend): match current linear API and test capability isolation - [`8cb9fda`](https://github.com/Comfy-Org/comfy-kitchen/commit/8cb9fdac282aee8f5bcdcb444a5f4000e17bbd27) docs(ascend): explain consolidated capability requirements ### 📊 Changes **6 files changed** (+1957 additions, -47 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+48 -12) 📝 `comfy_kitchen/backends/ascend/__init__.py` (+643 -26) ➕ `comfy_kitchen/backends/ascend/rope.py` (+365 -0) 📝 `tests/test_ascend_backend.py` (+601 -1) ➕ `tests/test_ascend_rope.py` (+292 -0) 📝 `tests/test_backends.py` (+8 -8) </details> ### 📄 Description ## Summary and relationship to the previous PRs Following the suggestion to consolidate the remaining Ascend work after #151 was merged, this PR starts directly from main (`d86e8d7`) and replaces the conflicting feature stack: | Previous PR | Included functionality | | --- | --- | | #153 | Dynamically quantized INT8 linear using `npu_quant_matmul` | | #154 | Interleaved/split-half RoPE and RMS-RoPE, single/paired and in-place APIs | | #155 | Optional fused RotateQuant for supported INT8 ConvRot calls | | #161 | Correctness-first ConvRot W4A4 integer-accumulation backend | The base backend from #151 is already in main and is **not reintroduced**. Feature implementations remain separate signed-off commits, followed by current-main compatibility fixes and documentation. W4A4 does not call the public INT8 linear API or require RotateQuant; RoPE is independent of the quantization features. ## Integration changes - Resolve the overlapping backend, capability table and test changes against main. - Preserve independent RoPE/RMS-RoPE capability detection, including operator-schema checks for `rotary_mode` and `div_mode`. - Preserve optional RotateQuant fallback, zero-scale handling and the K > 16000 fused-kernel guard. - Adapt INT8 linear to main's new `input_act_weight`, `input_act_eps`, `residual` and `residual_scale` arguments. Input RMSNorm and residual modulation use the shared device-side helpers, not a new fused GEMM implementation. - Test that missing quantization/matmul/RotateQuant capabilities do not disable unrelated functionality. - Preserve packed/strided RoPE views, in-place alias checks, partial RMS-RoPE and registry fallback for unsupported contracts. ## W4A4 correctness and performance scope This includes the **corrected** #161 implementation, not the earlier numerically different packed-kernel path: 1. Preserve FP32/FP16/BF16 activation dtype during rotation and use the reference signed-INT4 quantizer. 2. Unpack the same A4/W4 codes to INT8 and accumulate with `npu_quant_matmul` into INT32, using a unit scale and no per-token scale. 3. Cast and apply activation scale, weight scale and bias in eager's order; guard possible INT32 accumulator overflow. There is no compulsory FP32-to-BF16 input conversion or FP16 intermediate for FP32/BF16 outputs. FP16 retains the reference's own range limitations. This is not zero-copy packed-INT4 GEMM: unpacking adds memory/runtime cost. The earlier 134 s / 2.73x result from the superseded W4A4 implementation is **not a performance claim for this PR**. The integration workflow runs below are smoke tests, not repeated hot A/B measurements, and establish no new speedup or multi-seed quality guarantee. ## User-facing changes Eligible NPU calls can use the above Ascend implementations automatically through the existing registry. No mandatory MindIE-SD dependency, environment variable or public API change is introduced by this PR. Non-NPU import remains optional; unsupported contracts continue through the registry's existing fallback behavior. ## Validation One Ascend 910B3, physical card 4; PyTorch 2.10.0+cpu, torch_npu 2.10.0.post4, CANN 9.1.0. - Relevant NPU/backend/INT8/ConvRot/RoPE/constraint regression: **277 passed, 172 skipped** (two framework warnings). Skips cover other hardware/backend configurations. - CPU optional-import and backend/constraint/RoPE/RMS-RoPE regression with `TORCH_DEVICE_BACKEND_AUTOLOAD=0` and `sys.modules["torch_npu"] = None`: **361 passed, 598 skipped**, one warning. This explicitly simulates torch-npu being unavailable without uninstalling it from the shared validation container. - Replayed three saved real Krea activation/weight samples (from layers with K=6144 and K=16384): the integrated W4A4 implementation matches eager with `rtol=0, atol=0` in every case. - Fresh single-card ComfyUI smoke tests with the integrated checkout and node cache disabled: Z-Image INT8 and BF16 (1920x1088, 8 steps), Wan2.2 14B (640x640, 81 frames, 4 steps), and Krea2 mixed INT4 (1024x1024, seed 20260907, 8 steps) all completed successfully. Dispatch tracing confirmed Ascend INT8 linear/RMS-RoPE, Wan RoPE and Krea W4A4. Unsupported Krea RoPE calls retained eager fallback. - The new Krea image is pixel-identical (MAE=0) to the **saved** eager baseline from the #161 correctness audit; this was not a new two-arm A/B benchmark. All 81 output Wan frames decoded successfully. Z-Image/Krea images were also visually inspected. - `ruff check`, `ruff format --check` on all changed Python files and `git diff --check`: passed. ```bash ASCEND_RT_VISIBLE_DEVICES=4 OMP_NUM_THREADS=4 PYTHONPATH="$PWD" \ python -m pytest -q tests/test_ascend_backend.py tests/test_ascend_rope.py \ tests/test_backends.py tests/test_convrot_w4a4.py tests/test_int8.py \ tests/test_int8_input_act.py tests/test_constraints.py ``` The ComfyUI validation checkout retains the existing local custom-node shim routing mixed INT8 layers through Comfy Kitchen. That shim is not included in this PR. Workflow smoke results must not be interpreted as unmodified-ComfyUI support for every third-party custom node. ### Completed workflows, exact weights and output specifications The following are the exact checkpoint filenames selected in the saved successful API requests on 2026-09-25. All runs used one Ascend 910B3 (physical card 4), batch size 1, `--cache-none --disable-pinned-memory`, and `OMP_NUM_THREADS=4`. | Workflow | Main diffusion checkpoint(s) | Output (width × height) | Sampling | Observed request-to-completion time | | --- | --- | --- | --- | --- | | Z-Image Turbo INT8 ConvRot | `z_image_turbo_int8_convrot.safetensors` | One image, 1920 × 1088 | 8 steps; CFG 1; res_multistep/simple | 28.98 s | | Z-Image Turbo BF16 | `z_image_turbo_bf16.safetensors` | One image, 1920 × 1088 | 8 steps; CFG 1; res_multistep/simple | 23.76 s | | Wan2.2 T2V 14B FP16 + LightX2V | `wan2.2_t2v_high_noise_14B_fp16.safetensors` + `wan2.2_t2v_low_noise_14B_fp16.safetensors` | One video, 640 × 640; 81 frames at 16 FPS; **5.0625 s** | 4 steps total (2 high-noise + 2 low-noise); CFG 1; Euler/simple | 137.11 s | | Krea2 Turbo mixed INT4 | `krea2_turbo_int4_mixed.safetensors` | One image, 1024 × 1024 | 8 steps; CFG 1; Euler/simple | 356.64 s | Supporting weights used by these workflows: | Workflow | Text encoder | VAE | LoRA | | --- | --- | --- | --- | | Both Z-Image workflows | `qwen_3_4b.safetensors` | `ae.safetensors` | None | | Wan2.2 | `umt5_xxl_fp16.safetensors` | `wan_2.1_vae.safetensors` | `wan2.2_t2v_lightx2v_4steps_lora_v1.1_high_noise.safetensors` and `wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors`, each at strength 1.0 on its corresponding model | | Krea2 | `qwen3vl_4b_bf16.safetensors` | `qwen_image_vae.safetensors` | None | Seeds: both Z-Image runs used `493240645134828`; Wan used `832596732254118` for initial noise (the low-noise stage did not add noise); Krea used `20260907`. Krea used the existing `ComfyUI-INT4-Fast` custom node (`OTUNetLoaderW4A4`) and the local mixed-INT8 integration shim described above. Timing definition and limitations: - These are **single successful smoke-run wall times**, measured immediately before submitting the API request until successful history completion was observed. They include workflow execution, any model loading/transfers performed by that request, VAE decoding, output saving/video encoding, and polling overhead (approximately one second polling interval); they exclude service startup. - The workflows ran sequentially in the same service in the table order. Node output caching was disabled, but model residency/allocator state can carry across requests. Other validation checks also ran during the session. These are not isolated, uniformly cold or repeated hot measurements. - Therefore the table is execution evidence, **not an INT8-vs-BF16 comparison, an optimization speedup claim, or a performance guarantee**. - Image dimensions were checked from the generated files. The Wan MP4 metadata reports 640 × 640, 81 frames, 16 FPS and 5.0625 seconds, and all 81 frames decoded successfully. Video playback duration is separate from the 137.11-second generation wall time. --- <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:39 +00:00