CPU→GPU H2D copy in SymmetricPatchifier blocks CUDA graph capture every step
- 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
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