Comfy-Org / Comfy-Org/ComfyUI

CPU→GPU H2D copy in SymmetricPatchifier blocks CUDA graph capture every step

Open
#15,550 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 6h
Merged PRs (30d)
155

Description

## Summary

`SymmetricPatchifier.get_latent_coords()` in `comfy/ldm/lightricks/symmetric_patchifier.py:71-94` and the helper `latent_to_pixel_coords()` at `:9-33` both construct small tensors from Python scalars directly onto the GPU inside the per-step forward path. The resulting H2D copy from a non-pinned host buffer makes CUDA graph capture fail with `cudaErrorStreamCaptureUnsupported` / `cudaErrorStreamCaptureInvalidated`, and adds a small CPU-side launch cost on every step even when not capturing.

This blocks PyTorch CUDA graphs (`torch.cuda.graph`, `torch.compile(mode="reduce-overhead")`) on every LTXV/LTXAV workflow, and contributes measurably to per-step launch overhead.

## Repro

1. ComfyUI v0.32.0, any LTXV or LTXAV video workflow, `pip install torch>=2.5`
2. Apply around `model.apply_model`:

```python
import torch, types
im = model.inner_model # BaseModel subclass (LTXV)
target_step = 2
state = {"calls": 0}

def patched(self, x, t, *a, **kw):
state["calls"] += 1
if state["calls"] < target_step:
return orig(self, x, t, *a, **kw)
static_x, static_t = x.clone(), t.clone()
torch.cuda.synchronize()
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
out = orig(self, static_x, static_t, *a, **kw)
return out

orig = im.apply_model
im.apply_model = types.MethodType(patched, im)
```

3. Run a sampler step. Capture fails with:

```
torch.AcceleratorError: CUDA error: operation not permitted when stream is capturing
```

## Root cause

`comfy/ldm/lightricks/symmetric_patchifier.py::get_latent_coords` line 78:

```python
delta = torch.tensor(self._patch_size,
device=latent_sample_coords_start.device,
dtype=latent_sample_coords_start.dtype)[:, None, None, None]
```

This constructs a tensor from a Python tuple, placing it directly on the GPU. That is an H2D copy from a non-pinned buffer — illegal during graph capture, slow even outside capture.

A second instance of the same pattern exists at `latent_to_pixel_coords()` line 28:

```python
pixel_coords = (
latent_coords
* torch.tensor(scale_factors, device=latent_coords.device).view(*shape)
)
```

Both fire **every sampler step** because `get_latent_coords` is called via `process_timestep = model_base.py:1226` (LTXAV) on every forward.

## Suggested fix

Cache these small per-(patch_size, device, dtype) tensors in module-level dicts. The values are deterministic functions of `(patch_size, device, dtype)` and `(scale_factors, device, dtype)` respectively; nothing about them changes during a single sampler run.

Pseudocode:

```python
_PATCHIFIER_DELTA_CACHE = {}
_SCALE_FACTOR_CACHE = {}

# in get_latent_coords:
cache_key = (self._patch_size, latent_sample_coords_start.device,
latent_sample_coords_start.dtype)
delta = _PATCHIFIER_DELTA_CACHE.get(cache_key)
if delta is None:
delta = torch.tensor(self._patch_size, ...).to(...).view(...)
_PATCHIFIER_DELTA_CACHE[cache_key] = delta

# in latent_to_pixel_coords:
cache_key = (tuple(scale_factors), latent_coords.device, latent_coords.dtype, ...)
scale_t = _SCALE_FACTOR_CACHE.get(cache_key)
if scale_t is None:
scale_t = torch.tensor(scale_factors, device=...).view(*shape)
_SCALE_FACTOR_CACHE[cache_key] = scale_t
```

After the first step, both runs hit the cache and no H2D copies fire. CUDA graph capture proceeds past this point (next blocker is in the per-step `process_timestep` GPU-scalar reads, filed separately).

## Environment

- ComfyUI 0.32.0, python 3.12.11, torch 2.13.0+cu132
- NVIDIA RTX A6000 (46GB), driver-level `cudaMallocAsync`
- Observed while setting up CUDA graph capture / reduced-overhead sampling on LTXV-2.5 and LTXAV workflows

## Related

A parallel issue covers per-step GPU-to-host scalar reads in `comfy/ldm/minimax/model.py::_forward`. Both need to be fixed for full CUDA graph capture on video DiT models.

Contributor guide

Open the contributing guide

Research direction

Start with comfy/ldm/lightricks/symmetric_patchifier.py, especially latent_to_pixel_coords() and SymmetricPatchifier.get_latent_coords(), then trace their per-step use through model_base.py:1226. Reproduce the CUDA graph capture failure with the supplied LTXV or LTXAV workflow, cache the device- and dtype-specific tensors, and verify capture proceeds past these copies without recreating them each step.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.