Stochastic FP8 rounding regresses badly when the comfy-kitchen CUDA backend is disabled (cu128): LoRA + fp8_scaled model load never completes
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## Summary
On a CUDA 12.8 build, applying LoRAs to an fp8_scaled model that needs partial offload never finishes. It does not error, and it does not deadlock — it grinds indefinitely.
The regression comes from `comfy-kitchen` gaining `stochastic_rounding_fp8` (absent in 0.2.8, present in 0.2.26). Once it exists, `comfy/float.py` switches to it and abandons ComfyUI's own `manual_stochastic_round_to_float8`. On cu128 the comfy-kitchen **cuda** backend is disabled by `comfy/quant_ops.py`, and the **triton** backend does not implement this op, so the call lands on the **eager** (pure-PyTorch) reference implementation.
The two implementations are mathematically the same. The difference is that ComfyUI's version **chunks** the tensor and the comfy-kitchen path does not — and `comfy/float.py` additionally allocates a full-size uint8 RNG tensor up front.
So on cu128 the upgrade buys no kernel acceleration at all, while paying the cost of losing the chunking. With VRAM already near capacity during partial load, that is enough to make the weight-patching loop effectively never terminate.
## Environment
- ComfyUI v0.30.1
- Windows 10, RTX 4090 24 GB, 128 GB RAM
- torch 2.8.0.dev20250419+**cu128**
- comfy-kitchen 0.2.26, comfy-aimdo 0.4.11
- Launch flags: `--disable-dynamic-vram --reserve-vram 4 --use-sage-attention --fast fp16_accumulation`
- Model: Qwen Image Edit 2509, 19 GB fp8_scaled UNet, 5 LoRAs, native `UNETLoader`
## Symptom
Log stops after `Requested to load QwenImage` and never prints `loaded partially`. No exception, no OOM.
GPU telemetry during the stall, 8 samples 15 s apart, all identical in character:
```
utilization.gpu 100 % utilization.memory 0 % 23860-23934 MiB used ~106 W 47 C
```
100% "busy" with **zero memory-controller traffic** and no thermal rise. A 4090 doing real work draws 300-450 W. Left running long enough, the whole desktop becomes unresponsive.
## It is not a deadlock
`py-spy dump --nonblocking`, 8 samples over 2 minutes. The stack keeps moving:
```
samples 1-3:
synchronize (torch/cuda/__init__.py:1051)
load (comfy/model_patcher.py:1069) <- per-module device barrier
partially_load (comfy/model_patcher.py:1269)
model_use_more_vram (comfy/model_management.py:810)
samples 4-8:
cast_to (comfy/model_management.py:1535) <- r.copy_(weight, ...)
cast_to_device (comfy/model_management.py:1540)
calculate_weight (comfy/weight_adapter/lora.py:239 -> 236) <- line number moves
calculate_weight (comfy/lora.py:460)
patch_weight_to_device (comfy/model_patcher.py:909)
load (comfy/model_patcher.py:1067)
```
It is progressing, just orders of magnitude too slowly.
## The code on that path is byte-identical between v0.26.0 and v0.30.1
Every frame above was checked against `git diff v0.26.0 v0.30.1`: `cast_to`, `cast_to_device`, `calculate_weight` (both `comfy/lora.py` and `comfy/weight_adapter/` have zero diff), `patch_weight_to_device`, `ModelPatcher.load`, `ModelPatcher.partially_load`, `model_use_more_vram`, `model_load` — **none of them fall inside a diff hunk**. Of the 117 changed lines in `comfy/model_management.py`, the only ones on this path are two `detail()` log calls.
That is what pointed at the package rather than at core.
## Root cause
`comfy/float.py`:
```python
if _CK_STOCHASTIC_ROUNDING_AVAILABLE:
rng = torch.randint(0, 256, value.size(), dtype=torch.uint8, ..., generator=generator)
return _ck_stochastic_rounding_fp8(value, rng, dtype) # whole tensor, one shot
output = torch.empty_like(value, dtype=dtype) # chunked path below
num_slices = max(1, (value.numel() / (4096 * 4096)))
slice_size = max(1, round(value.shape[0] / num_slices))
for i in range(0, value.shape[0], slice_size):
output[i:i+slice_size].copy_(
manual_stochastic_round_to_float8(value[i:i+slice_size], dtype, generator=generator))
return output
```
`comfy_kitchen/backends/eager/quantization.py::stochastic_rounding_fp8` is essentially line-for-line the same algorithm as `manual_stochastic_round_to_float8`, but it receives the entire tensor, so the whole weight plus ~6 same-sized intermediates (`sign`, `abs_x`, `exponent`, `normal_mask`, `mantissa_scaled`, and the RNG tensor) are live simultaneously. The ComfyUI implementation bounds that to <= 16.7M elements per step.
Backend availability on this machine, from the startup log:
| backend | `stochastic_rounding_fp8` | status |
|---|---|---|
| cuda | yes | **disabled** by `comfy/quant_ops.py` (`cuda_version < 13`) |
| triton | no | n/a for this op |
| eager | yes | **actually used** |
Under comfy-kitchen 0.2.8 the log said `comfy_kitchen does not support stochastic FP8 rounding, please update comfy_kitchen.` and the chunked ComfyUI implementation was used. Same machine, same workflow, same flags: **80.92 s end to end**.
## Why almost nobody hits this
It requires the comfy-kitchen cuda backend to be disabled, i.e. torch built against CUDA < 13. On cu130+ the real kernel is used and there is no problem. Combined with needing fp8_scaled weights + LoRAs + enough VRAM pressure to force partial offload, the affected population is small — which is presumably why this has not been reported.
## Suggested direction (not a patch request, just where I would look)
Move the chunking loop outside the branch so both implementations are bounded, rather than only the fallback. That keeps the kernel for everyone who can use it and only removes the unbounded peak.
Caveat worth deciding on: chunking changes how RNG is consumed, so results would no longer be bit-identical to the current comfy-kitchen path for a given seed (statistically equivalent, but not reproducible against previous versions).
Alternatively this may belong in comfy-kitchen — an eager *reference* implementation arguably should not be handed an entire multi-GB tensor.
## Local workaround (for anyone who lands here)
Force `_CK_STOCHASTIC_ROUNDING_AVAILABLE = False` in `comfy/float.py`. This restores the chunked path. Verified: every number in the resulting `loaded partially` line matches the v0.26.0 baseline exactly — `16969.44 MB usable, 16898.61 MB loaded, 2585.34 MB offloaded, lowvram patches: 285`.
## Not the same as #15286
I tested that issue's workaround (forcing `read_tensor_file_slice_into` to return False, i.e. loading through RAM instead of the aimdo direct-read paths) and it changes nothing here. That report involves a **bf16** model, and `comfy/float.py::stochastic_rounding` returns early for fp32/fp16/bf16 — it never reaches the fp8 branch that breaks here. Same subsystem owner, different bug.
Contributor guide
Assessment
This issue has not been assessed yet.