[PR #52] Fix crash when running quantized models through torch.compile #107

New issue

Open

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

bsakaguchi commented

2026-10-26 22:20:27 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/52
Author: @jtreminio
Created: 6/25/2026
Status: 🔄 Open

Base: main ← Head: fix-quantized-tensor-torch-compile

📝 Commits (3)

54d5cac Fix crash when running quantized models through torch.compile
9789130 Merge branch 'main' into fix-quantized-tensor-torch-compile
1f3910a Make CodeRabbit happy

📊 Changes

2 files changed (+186 additions, -0 deletions)

View changed files

📝 comfy_kitchen/tensor/base.py (+15 -0)
➕ tests/test_unflatten_dynamic_shapes.py (+171 -0)

📄 Description

Quantized models could crash during torch.compile whenever the input size changes between runs, for example an image model where the prompt length varies, because the quantized tensor wasn't reporting its shape the way PyTorch's compiler expects when it rebuilds the tensor.

This change makes that step pass the requested shape and stride through correctly, so compiling quantized models (including mxfp8) now works instead of erroring out, with no change to normal non-compiled use. Includes a regression test that reproduces the original failure.

An example model that crashed is https://huggingface.co/Comfy-Org/Krea-2/blob/main/diffusion_models/krea2_turbo_mxfp8.safetensors

Full disclosure: I am a software engineer but Torch internals are a mystery to me. This PR was heavily assisted with Claude, but I did the code review, I did the testing, I verified this works as intended, and minimized the area of impact for the fix.

_{🔄 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/52 **Author:** [@jtreminio](https://github.com/jtreminio) **Created:** 6/25/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix-quantized-tensor-torch-compile` --- ### 📝 Commits (3) - [`54d5cac`](https://github.com/Comfy-Org/comfy-kitchen/commit/54d5cacdfb3555e0f4802bd9949af01fe4725621) Fix crash when running quantized models through torch.compile - [`9789130`](https://github.com/Comfy-Org/comfy-kitchen/commit/978913023b9f04046a2ecfb8bc0ebaaff6eb872f) Merge branch 'main' into fix-quantized-tensor-torch-compile - [`1f3910a`](https://github.com/Comfy-Org/comfy-kitchen/commit/1f3910a1e26bd321cb5d3e01995bc324494ce773) Make CodeRabbit happy ### 📊 Changes **2 files changed** (+186 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/tensor/base.py` (+15 -0) ➕ `tests/test_unflatten_dynamic_shapes.py` (+171 -0) </details> ### 📄 Description Quantized models could crash during `torch.compile` whenever the input size changes between runs, for example an image model where the prompt length varies, because the quantized tensor wasn't reporting its shape the way PyTorch's compiler expects when it rebuilds the tensor. This change makes that step pass the requested shape and stride through correctly, so compiling quantized models (including mxfp8) now works instead of erroring out, with no change to normal non-compiled use. Includes a regression test that reproduces the original failure. An example model that crashed is https://huggingface.co/Comfy-Org/Krea-2/blob/main/diffusion_models/krea2_turbo_mxfp8.safetensors --- Full disclosure: I am a software engineer but Torch internals are a mystery to me. This PR was heavily assisted with Claude, but I did the code review, I did the testing, I verified this works as intended, and minimized the area of impact for the fix. --- <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:20:27 +00:00

bsakaguchi referenced this issue

2026-10-26 22:20:52 +00:00

[PR #107] Add MPS fallback for INT8 linear #148

bsakaguchi referenced this issue

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

[PR #145] Add a float-GEMM fallback for int8_linear on devices without aten::_int_mm #173

bsakaguchi referenced this issue

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