Refresh the whole captured input before a CUDA graph replay - #8488
Refresh the whole captured input before a CUDA graph replay#8488vineethsaivs wants to merge 1 commit into
Conversation
A replay runs the tensors that were captured into the graph, so a later call's
values reach the model only by being copied into them. Every _graph_replay
matched on torch.is_tensor and therefore stopped at the top level, leaving any
argument that keeps its tensors inside a container holding what capture saw.
SDXL is the clearest case. Its pipeline passes
added_cond_kwargs={"text_embeds": ..., "time_ids": ...}
to the UNet, which is a dict and not a tensor, so it was never refreshed. Every
image after the first was denoised with the first prompt's conditioning, with no
error to show for it.
Add refresh_static_tensors() next to CUDAGraph and use it at all six replay
sites: DSUNet, DSVAE's decoder, encoder and combined graphs, DSClipEncoder and
InferenceEngine. It walks dicts, lists and tuples and copies at the tensors.
The loops now zip over the captured side instead of indexing by the new call's
length, so an extra positional argument or an unseen kwarg leaves the captured
values alone rather than raising IndexError or KeyError from inside the replay
path.
Test: tests/unit/inference/test_cuda_graph_static_inputs.py covers the walk over
dicts, lists and tuples, the cases it must leave alone, and DSUNet._graph_replay
refreshing a nested added_cond_kwargs. The last one fails on master, where
text_embeds keeps its capture-time value. CPU only, no graph capture needed,
since the bug is entirely in which tensors get copied.
Signed-off-by: Vineeth Sai <vineethsai4444@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d8f9abcb5b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for key, captured in self.static_kwargs.items(): | ||
| if key in kwargs: | ||
| refresh_static_tensors(captured, kwargs[key]) |
There was a problem hiding this comment.
When an SDXL call supplies added_cond_kwargs, this replay code does refresh the captured dictionary, but that dictionary is never consumed: _create_cuda_graph invokes DSUNet._forward, and _forward accepts added_cond_kwargs yet only forwards cross_attention_kwargs (lines 65-80). Thus the new copies cannot affect either capture or replay, so SDXL conditioning remains absent rather than being updated for later prompts. The new test calls _graph_replay directly and consequently misses this end-to-end path.
Useful? React with 👍 / 👎.
1 participant
What broke
_graph_replaycopies a later call's values into the tensors the CUDA graph captured, but it only matchedtorch.is_tensor, so it stopped at the top level. An argument that keeps its tensors inside a container was never refreshed.SDXL is the clearest case. Its pipeline calls the UNet with
That is a dict, not a tensor, so it was skipped. Every image after the first was denoised with the first prompt's pooled embeddings and time ids. The pipeline still returns a plausible image, so nothing reports an error.
Why
The same six lines are repeated at every replay site, and all six have the same gap:
DSUNet._graph_replayDSVAE._graph_replay_decoder,_graph_replay_encoder,_graph_replayDSClipEncoder._graph_replayInferenceEngine._graph_replayWhat changed
refresh_static_tensors()lands next toCUDAGraphand all six sites call it. It walks dicts, lists and tuples and copies at the tensors.The loops now zip over the captured side instead of indexing by the new call's length. An extra positional argument or a kwarg that was not present at capture now leaves the captured values alone, rather than raising
IndexErrororKeyErrorfrom inside the replay path.Tests
tests/unit/inference/test_cuda_graph_static_inputs.py:DSUNet._graph_replayrefreshing a nestedadded_cond_kwargsThe last one fails on master, where
text_embedskeeps its capture-time value, and passes here. CPU only and no graph capture, because the bug is entirely in which tensors get copied.yapf0.40.0 andflake8with the repo config are clean on all six files.