[PR #253] Route INT8 linear_input_act through the int8_linear custom op #258

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/253
Author: @qweqweewqe7-create
Created: 10/17/2026
Status: 🔄 Open

Base: main ← Head: moss-int8-input-act-compile

📝 Commits (1)

6a9c452 Route INT8 linear_input_act through the int8_linear custom op

📊 Changes

1 file changed (+6 additions, -6 deletions)

View changed files

📝 comfy_kitchen/tensor/int8.py (+6 -6)

📄 Description

Summary

TensorWiseINT8Layout.linear_input_act called ck.int8_linear directly. Its backend implementation hands raw pointers to the compiled extension, so torch.compile traced into it and failed on a FakeTensor __dlpack__:

torch._dynamo.exc.TorchRuntimeError: RuntimeError when making fake tensor call
  call_method __dlpack__(*(FakeTensor(..., size=(6187, 16384), dtype=torch.bfloat16),), **{'stream': -1})
  ...
  File "comfy_kitchen/tensor/int8.py", in linear_input_act
    return ck.int8_linear(

This shows up with ComfyUI's native MOSS-Image model, whose ffn_down uses linear_input_act(..., "situ"), under a per-block or whole-model torch.compile.

linear_input_act now calls torch.ops.comfy_kitchen.int8_linear, which already takes input_act, input_act_weight, input_act_eps, residual and residual_scale, the same way the plain linear path does. Eager behaviour and results are unchanged.

Validation

RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), PyTorch 2.16.0a0+rocm10.2.0a20261015:

MOSS-Image 1024x1536 DiT with torch.compile(block, fullgraph=True) per block (RX 9070 XT): fails without this change, passes with it.
tests/test_int8.py tests/test_int8_input_act.py tests/test_int8_residual.py and the HIP/ConvRot/adaln/rope test files on both GPUs: same results as unmodified main. The only failures, 18 triton cases and test_int8_linear_correctness[hip], fail identically on main itself.

Update: rebased onto v0.2.40

Rebased onto main (f8b7e97, v0.2.40). The only conflict was the import block. import comfy_kitchen as ck is still unused after this change, and #259's NoCapableBackendError import is kept.

Reproduction on RX 9070 XT (gfx1201) with PyTorch 2.16.0a0+rocm10.3.0a20261023:

Setup: torch.compile(fullgraph=True) over TensorWiseINT8Layout.linear_input_act(..., input_act="situ"). The weight is a per-channel ConvRot 3072 x 8192, and x is 512 x 16384.
main still fails, now on a traced logging.Logger call inside ck.int8_linear.
This branch compiles, and its output is equal to eager.

Tests on RX 9070 XT and RX 9060 XT (test_hip_wmma, test_hip_dispatch, test_gemm_bias_contract, test_int8*, test_adaln, test_rms_rope, test_int8_attention, test_sol_attn, test_convrot_unpacked_quantization, test_w4a8_mma, test_w6a8) show the same failures as unmodified main: 18 triton rope cases and test_int8_linear_correctness[hip].

🤖 Generated with Claude Code

_{🔄 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/253 **Author:** [@qweqweewqe7-create](https://github.com/qweqweewqe7-create) **Created:** 10/17/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `moss-int8-input-act-compile` --- ### 📝 Commits (1) - [`6a9c452`](https://github.com/Comfy-Org/comfy-kitchen/commit/6a9c452facc4042eec1da481922e10693319bfd7) Route INT8 linear_input_act through the int8_linear custom op ### 📊 Changes **1 file changed** (+6 additions, -6 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/tensor/int8.py` (+6 -6) </details> ### 📄 Description ## Summary `TensorWiseINT8Layout.linear_input_act` called `ck.int8_linear` directly. Its backend implementation hands raw pointers to the compiled extension, so `torch.compile` traced into it and failed on a FakeTensor `__dlpack__`: ``` torch._dynamo.exc.TorchRuntimeError: RuntimeError when making fake tensor call call_method __dlpack__(*(FakeTensor(..., size=(6187, 16384), dtype=torch.bfloat16),), **{'stream': -1}) ... File "comfy_kitchen/tensor/int8.py", in linear_input_act return ck.int8_linear( ``` This shows up with ComfyUI's native MOSS-Image model, whose `ffn_down` uses `linear_input_act(..., "situ")`, under a per-block or whole-model `torch.compile`. `linear_input_act` now calls `torch.ops.comfy_kitchen.int8_linear`, which already takes `input_act`, `input_act_weight`, `input_act_eps`, `residual` and `residual_scale`, the same way the plain linear path does. Eager behaviour and results are unchanged. ## Validation RX 9070 XT (gfx1201) and RX 9060 XT (gfx1200), PyTorch `2.16.0a0+rocm10.2.0a20261015`: - MOSS-Image 1024x1536 DiT with `torch.compile(block, fullgraph=True)` per block (RX 9070 XT): fails without this change, passes with it. - `tests/test_int8.py tests/test_int8_input_act.py tests/test_int8_residual.py` and the HIP/ConvRot/adaln/rope test files on both GPUs: same results as unmodified `main`. The only failures, 18 triton cases and `test_int8_linear_correctness[hip]`, fail identically on `main` itself. ## Update: rebased onto v0.2.40 Rebased onto `main` (`f8b7e97`, v0.2.40). The only conflict was the import block. `import comfy_kitchen as ck` is still unused after this change, and #259's `NoCapableBackendError` import is kept. Reproduction on RX 9070 XT (gfx1201) with PyTorch `2.16.0a0+rocm10.3.0a20261023`: - Setup: `torch.compile(fullgraph=True)` over `TensorWiseINT8Layout.linear_input_act(..., input_act="situ")`. The weight is a per-channel ConvRot 3072 x 8192, and `x` is 512 x 16384. - `main` still fails, now on a traced `logging.Logger` call inside `ck.int8_linear`. - This branch compiles, and its output is equal to eager. Tests on RX 9070 XT and RX 9060 XT (`test_hip_wmma`, `test_hip_dispatch`, `test_gemm_bias_contract`, `test_int8*`, `test_adaln`, `test_rms_rope`, `test_int8_attention`, `test_sol_attn`, `test_convrot_unpacked_quantization`, `test_w4a8_mma`, `test_w6a8`) show the same failures as unmodified `main`: 18 triton rope cases and `test_int8_linear_correctness[hip]`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <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:22:11 +00:00

bsakaguchi referenced this issue

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

[PR #272] Decline fp16-accumulate CUTLASS kernels on sub-sm_80 GPUs #278