Comfy-Org / Comfy-Org/ComfyUI

MiniMax H3 _forward does per-step GPU→CPU syncs (sigma_v, text_tags) — graph capture blocked

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

Description

## Summary

`comfy/ldm/minimax/model.py::MiniMaxH3Model._forward` performs multiple per-step device-to-host scalar reads and tensor-to-Python conversions that force CUDA sync and break CUDA graph capture / `torch.compile(mode="reduce-overhead")`.

| Location | Code | Kind |
|----------|------|------|
| `:538` | `sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6)` then read as scalar at `:541` | D2H of GPU scalar |
| `:540` (original) | `t_a = float(1.0 - time_shift_sigma(...))` | D2H |
| `:561` (original) | `tags = text_tags.view(-1).tolist()` | D2H |

`float(gpu_scalar)` or `tensor.tolist()` inside the forward path stalls the host on the in-flight GPU stream and invalidates stream capture with `cudaErrorStreamCaptureUnsupported`.

## Why this is hard to fully fix

Beyond the three obvious syncs, the per-step layout machinery iterates `layout.segments` in Python, indexes into `t_row[seg_t[kind]]`, and walks text-tag runs before attending. Each downstream decision depends on the scalar `sigma_v`, so the entire segment-dispatch tree is data-dependent Python control flow rooted in a GPU tensor read. Mitigating the surface syncs is possible (done downstream by caching the text-token list; sigma arithmetic inlined on host), but a true graph capture of `_forward` requires hoisting the sigma-to-layout mapping out of the model — e.g., precomputing per-sigma segments and passing a step index through `minimax_payload`, or reshaping the segment logic to be tensor-driven.

## Concrete minimal patch (partial)

```python
# was: t_v = float(1.0 - sigma_v); t_a = float(1.0 - time_shift_sigma(...))
# now: one host pull, the rest in pure Python floats
sigma_v_f = float(sigma_v)
t_v = 1.0 - sigma_v_f
_base = sigma_v_f / (shift_v + sigma_v_f * (1.0 - shift_v))
t_a = 1.0 - shift_a * _base / (1.0 + (shift_a - 1.0) * _base)
```

And for the text tags (run once, cached on the payload to skip the per-step D2H):

```python
if text_tags is not None and "_text_token_tags_list" not in payload:
payload["_text_token_tags_list"] = text_tags.view(-1).tolist()
tags = payload.get("_text_token_tags_list")
```

These two reduce per-step D2H events from three to one, but graph capture still fails because the Python layout machinery fundamentally depends on sigma as a scalar.

## What would unblock full graph capture

Pass an integer step index into the model and precompute the sigma-indexed layout decisions once at sampler init. The model forward then becomes a pure function of tensor inputs with static control flow for a given graph instance.

## Environment

- ComfyUI 0.32.0, torch 2.13.0+cu132, python 3.12.11
- RTX A6000 (sm86), `cudaMallocAsync`
- MiniMax H3 ref2va, `Hq=32, Sq=80661, D=128`, bf16, no offload

## Related

Sibling issue for the LTX-AV `SymmetricPatchifier` H2D copy: Comfy-Org/ComfyUI#15550.

Contributor guide

Open the contributing guide

Research direction

Start in comfy/ldm/minimax/model.py at MiniMaxH3Model._forward and trace the sigma_v, time-shift, text_tags, and layout.segments paths described in the issue. Compare the partial caching and host-arithmetic mitigation with the requirements for CUDA graph capture and torch.compile; done means eliminating the remaining data-dependent host synchronization and validating graph capture for this forward path.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.