Skip to content

vllm_omni.diffusion.models.auk.cudagraph_wrapper

Per-shape CUDA graph for one AuK DiT denoise step.

logger module-attribute

logger = init_logger(__name__)

AuKCUDAGraphWrapper

Replay one DiT denoise step and leave Euler scheduling to the caller.

The graph is keyed by the target, text and reference sequence lengths, whether the CFG branch is enabled and, when the caller passes the time grid, its length. Timestep (or the step index), and the CFG strength are mutable scalar buffers, so all Euler steps and CFG values within one path reuse one graph. With the time grid, the adaLN modulations of every step are part of the per-request context and the graph selects its row.

Only :meth:AuKTransformer.step is captured. The per-request context (text projection, reference embedding, padding biases, rotary tables) is built by :meth:AuKTransformer.prepare on the first step of a request and copied into the graph's static context, so the replayed steps skip it.

dit instance-attribute

dit = dit

enabled instance-attribute

enabled = bool(enabled)

max_graphs instance-attribute

max_graphs = max(1, int(max_graphs))