This function is used to handle context parallel behavior, split input tensors into multiple parts and scatter them to different GPUs. I

Decode the quantized tensor. Args: quant (torch.Tensor): Quantized tensor. Returns: torch.Tensor: D

Encode the input tensor. Args: video (torch.Tensor): Input tensor with shape (N, T, C, H, W). vae (VideoToken

Get the current EventPathTimer instance. Returns: EventPathTimer: The current EventPathTimer instance. Raises: AssertionErro

Non-Const Method

Forward pass of PixArt, but also batches the unconditional forward pass for classifier-free guidance.

Const Method

Const Method

Const Method

Generate orthogonal parallel groups based on the parallel size and mask. Arguments: world_size (int): world size parallel_size (

Const Method

When seq_len is very small, CPU-bound issues are severe. By fusing kv communication, CPU operations and the number of kernel launches

By fusing the communication of q, k, and v together, further optimize CPU-bound issues.

Forward embedding and meta for VideoDiT. NOTE: This function should only handle single card behavior. Input: x:

Init intervel

Initialize the PPScheduler instance. Raises: AssertionError: If the PPScheduler is already initialized.

Init Timestep and Transform t

Non-Const Method

Useful for code segments that may be accessed with or without mpu initialization

generate N-D grid in dimension order. The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.

generate N-D grid in dimension order. The ndgrid function is like meshgrid except that the order of the first two input arguments are switched.

Get a tensor from the receive queue and wait for the receive operation to complete. Returns: torch.Tensor: The received tensor ob

Inputs: x: (s, b, h). Tensor of inputs embedding (images or latent representations of images) c_mapping: (s, b). Tensor of condit

Receive data from the previous pipeline and return the received tensor. Args: shape (torch.Size): The shape of the tensor to rece

Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes fro

Split a tensor along its last dimension. Args: tensor: input tensor. num_partitions: number of partitions to split the tensor

Record the current time with a message. Args: message (str): A message to log along with the current time. This

Decodes the input tensor using the tile autoencoder. Args: x (torch.Tensor): The input tensor to be decoded.

Encodes the input tensor `x` using tiled encoding. Args: x (torch.Tensor shape:[N C T H W]): The input tensor to be enco

Create sinusoidal timestep embeddings. :param t: a 1-D Tensor of N indices, one per batch element. These ma

Non-Const Method