Fix crash when running quantized models through torch.compile

[jtreminio]

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.

[github-actions]

✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
_{Posted by the CLA Assistant Lite bot.}

[jtreminio]

I have read and agree to the Contributor License Agreement

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 88109be2-3b41-4ddb-8270-3c5480926943

📥 Commits

Reviewing files that changed from the base of the PR and between 9789130 and 1f3910a.

📒 Files selected for processing (1)

• tests/test_unflatten_dynamic_shapes.py

---

📝 Walkthrough

Walkthrough

QuantizedTensor now threads outer_stride through wrapper construction and unflatten rebuilds when outer_size is a real sequence. The PR also adds CPU, symbolic-shape, and CUDA torch.compile regressions for dynamic unflatten behavior.

Changes

Dynamic outer stride rebuilds

Layer / File(s)	Summary
Wrapper-subclass stride plumbing comfy_kitchen/tensor/base.py	QuantizedTensor.__new__ accepts outer_stride and passes it to _make_wrapper_subclass as strides=, and __init__ accepts the same argument.
Unflatten rebuild path comfy_kitchen/tensor/base.py	__tensor_unflatten__ replaces orig_shape for sequence outer_size values and forwards outer_stride into the rebuilt tensor.
Concrete CPU regressions tests/test_unflatten_dynamic_shapes.py	The test module adds helpers and CPU cases for explicit outer shape and stride rebuilds, placeholder handling, and eager real-layout roundtrips.
Symbolic and compiled regressions tests/test_unflatten_dynamic_shapes.py	The remaining tests cover FakeTensorMode symbolic outer shapes and a CUDA torch.compile path with dynamic batch sizes.

🚥 Pre-merge checks | ✅ 2

✅ Passed checks (2 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

✨ Simplify code

• Create PR with simplified code

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/test_unflatten_dynamic_shapes.py`:
- Around line 56-62: The `_dummy_qt` helper is registering the "Dummy" layout on
every call via register_layout_class, which repeatedly mutates the global
LAYOUTS registry across tests. Move the register_layout_class("Dummy",
self._DummyLayout) call out of _dummy_qt and into module scope so the layout is
registered once when the test module loads, while leaving _dummy_qt only
responsible for constructing the QuantizedTensor with _DummyLayout and
_DummyParams.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 38c59d67-74a6-44d9-9c76-3680e42b8ba5

📥 Commits

Reviewing files that changed from the base of the PR and between 4bd0f9b and 54d5cac.

📒 Files selected for processing (2)

• comfy_kitchen/tensor/base.py
• tests/test_unflatten_dynamic_shapes.py