CUDA-graph replay of fixed-shape keyword-only forwards.
Written for TaoMate-H3's audio teacher, which runs nine short, launch-bound DiT forwards per five-second request; replaying a captured graph removes the Python launch path. XPUGraphedForward (xpu_graph.py) derives from this class and overrides only the device hooks below; pick the class for the current platform with current_omni_platform.get_graphed_forward_cls().
GraphedForward wraps a keyword-only forward. Calls are keyed by the structure of their kwargs (tensor shapes/dtypes and every non-tensor value); the first call of a key captures a graph on static copies of the tensors and later calls copy the live tensors into those buffers and replay. Values that only change per request (prompt embeddings, RoPE table, row positions) are therefore ordinary inputs and never force a re-capture; only a new document shape does (a prompt with another token count, the first request's longer audio, a reference tail present or not).
Graph capture forbids host synchronization inside the captured region, so the wrapper (a) enters the caller's capture_context around the capture forwards and around eager forwards inside an outer capture (TaoMate uses it to bypass its host-keyed AdaLN projection cache), and (b) validates the very first capture with torch.cuda.set_sync_debug_mode("error") before capturing, so a synchronizing operator fails loudly on every rank at the same program point instead of leaving peers waiting in a collective. Any failure disables the wrapper for the rest of the process (agreed across the sequence-parallel group when one is given) and the caller falls back to eager execution.
Capture-and-replay wrapper around a keyword-only nn.Module forward.
The device-specific steps are the hooks _graphs_supported, _warmup_stream and _capture_graph (plus _sync_probe); a platform subclass overrides those and inherits keying, static buffers, replay, LRU eviction and rank agreement.
Deep-copy the tensors of a kwargs tree into fresh contiguous buffers.
Tensors are reached through dicts, lists and tuples only. Any other object (a layout dataclass, a device, a dtype) is passed through unchanged and must not hold a tensor: it would keep its capture-time address without being refreshed on replay, so such objects are rejected.
Copy the tensors of current into the matching buffers of static.
Returns the number of tensors copied. The trees must share one signature.