[GH-ISSUE #204] dequantize_int8_simple_dtype: cuda/hip drop leading dims when a one-element scale has more dims than q #60

Open
opened 2026-10-26 22:19:49 +00:00 by bsakaguchi · 0 comments
Owner

Originally created by @arcusbuilds on GitHub (Sep 26, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/204

cuda and hip use scalar mode whenever scale.numel() == 1. Eager does q.float() * scale, which broadcasts, so a one-element scale with more dims than q changes the output shape:

import torch, comfy_kitchen as ck

q = torch.randint(-128, 128, (3, 4), dtype=torch.int8, device="cuda")
scale = torch.tensor([[[0.01]]], device="cuda")  # shape (1, 1, 1)

with ck.registry.use_backend("eager"):
    print(ck.dequantize_int8_simple(q, scale).shape)  # torch.Size([1, 3, 4])
with ck.registry.use_backend("cuda"):
    print(ck.dequantize_int8_simple(q, scale).shape)  # torch.Size([3, 4])

Values match; only the shape differs. The check is in backends/cuda/__init__.py and backends/hip/__init__.py. ascend is fine. The Triton version in #197 copies the cuda check.

Fix: use scalar mode only when scale.dim() <= q.dim(), and send other shapes to the existing eager fallback. PR to follow. Real callers pass a scalar or a [rows, 1] scale, so this is low impact.

Originally created by @arcusbuilds on GitHub (Sep 26, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/204 cuda and hip use scalar mode whenever `scale.numel() == 1`. Eager does `q.float() * scale`, which broadcasts, so a one-element scale with more dims than `q` changes the output shape: ```python import torch, comfy_kitchen as ck q = torch.randint(-128, 128, (3, 4), dtype=torch.int8, device="cuda") scale = torch.tensor([[[0.01]]], device="cuda") # shape (1, 1, 1) with ck.registry.use_backend("eager"): print(ck.dequantize_int8_simple(q, scale).shape) # torch.Size([1, 3, 4]) with ck.registry.use_backend("cuda"): print(ck.dequantize_int8_simple(q, scale).shape) # torch.Size([3, 4]) ``` Values match; only the shape differs. The check is in `backends/cuda/__init__.py` and `backends/hip/__init__.py`. ascend is fine. The Triton version in #197 copies the cuda check. Fix: use scalar mode only when `scale.dim() <= q.dim()`, and send other shapes to the existing eager fallback. PR to follow. Real callers pass a scalar or a `[rows, 1]` scale, so this is low impact.
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
Comfy-Org/comfy-kitchen#60
No description provided.