Comfy-Org / Comfy-Org/comfy-kitchen
eager backend advertises int8_linear on MPS but dispatches to CUDA-only torch._int_mm
- Dominant language
- Python
- Stars
- 220
- Forks
- 91
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 12
Description
### Describe the bug
int8 quantized checkpoints fail immediately on Apple Silicon. The eager backend lists `int8_linear` under native ops on MPS, but `fast_int8_mm` calls `torch._int_mm`, which has no MPS kernel. It dies at the first linear layer of step 0.
Interestingly, the nvfp4 text encoder in the same workflow loads and runs fine. It's classified as emulated. So the emulation path works on MPS. The problem looks like just the native/emulated classification being wrong for int8 on this device.
### Reproduce
Files (from the MiniMax H3 template):
- `diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors`
- `text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors`
Load and run any sampler. Fails at 0/20.
Startup shows:
```
Found comfy_kitchen backend cuda: available=False (Extension file not found: .../backends/cuda/_C.abi3.so)
Found comfy_kitchen backend triton: available=False (No module named 'triton')
Found comfy_kitchen backend hip: available=False (PyTorch ROCm/HIP runtime not available)
Found comfy_kitchen backend eager: available=True
```
Model load shows:
```
Detected mixed precision quantization
Native ops: int8_tensorwise, convrot_w4a4 , emulated ops: float8_e4m3fn, mxfp8, nvfp4, float8_e5m2
```
### Traceback
```
comfy/ops.py:1313 in _forward
return torch.nn.functional.linear(input, weight, bias)
comfy_kitchen/tensor/base.py:362 in __torch_dispatch__
comfy_kitchen/tensor/int8.py:287 in _handle_int8_linear_tensorwise
comfy_kitchen/backends/eager/quantization.py:1222 in _op_int8_linear
comfy_kitchen/backends/eager/quantization.py:1022 in int8_linear
result = _int8_matmul_accumulate(x_8, weight.T.contiguous())
comfy_kitchen/backends/eager/quantization.py:779 in _int8_matmul_accumulate
result = fast_int8_mm(a, b)
comfy_kitchen/backends/eager/quantization.py:754 in fast_int8_mm
return torch._int_mm(lhs, rhs)
NotImplementedError: The operator 'aten::_int_mm' is not currently
implemented for the MPS device.
```
Upstream op tracker: [[pytorch/pytorch#141287](https://github.com/pytorch/pytorch/issues/141287)]
### Environment
- macOS 26.1, Apple Silicon, 48GB unified
- torch 2.10.0, Python 3.12.12
- comfy-kitchen 0.2.26, ComfyUI 0.30.0
- Device mps, vram state SHARED
### Notes
Not a regression. [[#9](https://github.com/Comfy-Org/comfy-kitchen/issues/9)](https://github.com/Comfy-Org/comfy-kitchen/issues/9) scoped the int8 path to Nvidia 8.0-8.9 with `torch._int_mm` as the matmul, so MPS was never considered. Just a gap.
The README says out-of-domain requests fall back to torch/eager, and describes HIP degrading on RDNA2 by not advertising GEMMs it can't run. Same idea seems like it should apply here. Eager is the fallback tier, so it shouldn't advertise something it can't service on the current device.
Two ways I can see to fix it:
**1. Don't advertise `int8_linear` on non-CUDA devices**, so dispatch falls through to the normal torch path. Probably the cleaner one given the HIP precedent.
**2. Fall back inside `fast_int8_mm`:**
```python
def fast_int8_mm(lhs, rhs):
if lhs.device.type != "cuda":
return torch.matmul(
lhs.to(torch.float32), rhs.to(torch.float32)
).to(torch.int32)
return torch._int_mm(lhs, rhs)
```
fp32 holds integers exactly to 2^24, so this should be bit-identical to `_int_mm` for realistic K with int8 inputs. Loses the int8 speedup but stays on GPU. (`PYTORCH_ENABLE_MPS_FALLBACK=1` "works" but round-trips to CPU on every linear in every block on every step, so it's not really usable.)
Happy to open a PR for either. Let me know which you'd prefer.
Contributor guide
Research direction
Start in comfy_kitchen/backends/eager/quantization.py, reading the eager native-op selection and the fast_int8_mm path around the traceback. Reproduce the int8 sampler failure on MPS, then verify that int8_linear no longer reaches the unsupported torch._int_mm path and the workflow completes its first step without CPU fallback.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 65/100