(self, decode_fn: Callable[..., torch.Tensor], *args: Any, **kwargs: Any)
| 110 | torch.cuda.synchronize(device) |
| 111 | |
| 112 | def submit(self, decode_fn: Callable[..., torch.Tensor], *args: Any, **kwargs: Any) -> None: |
| 113 | self._num_submitted += 1 |
| 114 | if not self.enabled: |
| 115 | self._sync_if_cuda(args) |
| 116 | t0 = time.perf_counter() |
| 117 | self._chunks.append(decode_fn(*args, **kwargs)) |
| 118 | self._sync_if_cuda(args) |
| 119 | chunk_ms = (time.perf_counter() - t0) * 1000.0 |
| 120 | self._sync_decode_ms += chunk_ms |
| 121 | logger.info( |
| 122 | "[AsyncVAEChunkDecoder] sync VAE chunk {}/{} decode={:.6f} seconds", |
| 123 | self._num_submitted, |
| 124 | self._num_submitted, |
| 125 | chunk_ms / 1000.0, |
| 126 | ) |
| 127 | return |
| 128 | |
| 129 | device = self._resolve_device(args) |
| 130 | stream = self._ensure_stream(device) |
| 131 | # LongLive: diffusion_done.record() then vae_stream.wait_event(diffusion_done) |
| 132 | diffusion_done = torch.cuda.Event() |
| 133 | diffusion_done.record(torch.cuda.current_stream(device)) |
| 134 | |
| 135 | # LongLive waits for the previous chunk's VAE before queueing the next |
| 136 | # decode (stateful cached VAE). Measure exposed wait here; do not block |
| 137 | # the next segment's DiT — that overlap happens after submit() returns. |
| 138 | if self._prev_vae_done is not None: |
| 139 | t0 = time.perf_counter() |
| 140 | self._prev_vae_done.synchronize() |
| 141 | wait_ms = (time.perf_counter() - t0) * 1000.0 |
| 142 | self._submit_wait_ms += wait_ms |
| 143 | self._log_new_async_chunks(wait_ms, wait_kind="submit_wait") |
| 144 | |
| 145 | with torch.cuda.stream(stream), torch.no_grad(): |
| 146 | stream.wait_event(diffusion_done) |
| 147 | decode_start = torch.cuda.Event(enable_timing=True) |
| 148 | decode_end = torch.cuda.Event(enable_timing=True) |
| 149 | decode_start.record(stream) |
| 150 | output = decode_fn(*args, **kwargs) |
| 151 | decode_end.record(stream) |
| 152 | self._chunks.append(output) |
| 153 | self._decode_events.append((decode_start, decode_end)) |
| 154 | self._prev_vae_done = torch.cuda.Event() |
| 155 | self._prev_vae_done.record(stream) |
| 156 | for value in list(args) + list(kwargs.values()): |
| 157 | if isinstance(value, torch.Tensor) and value.device.type == "cuda": |
| 158 | value.record_stream(stream) |
| 159 | if isinstance(output, torch.Tensor) and output.device.type == "cuda": |
| 160 | output.record_stream(stream) |
| 161 | |
| 162 | def finish(self) -> list[torch.Tensor]: |
| 163 | if self._stream is not None: |
no test coverage detected