Output from one generation block.

Configuration for :class:`ColaDiTModel`. Parameterizes the **block-causal Diffusion Transformer prior** ``p_psi(z_0)`` of Cola DLM (Eq. 2.1.4

Block-by-block generation engine for Cola-DLM. Usage:: engine = ColaEngine(dit, vae, tokenizer) for block_text in engine.generat

Small synchronous wrapper around the Cola DLM inference pipeline.

Configuration for :class:`ColaTextVAEModel`. Parameterizes the **Text VAE** of Cola DLM, which provides *both* the inference encoder ``q_phi(

Distributed Muon+AdamW. Handles gradient sync internally (no DDP wrapper needed). Uses 3-phase async communication: 1. Launch all reduce

Combined optimizer: Muon for 2D matrix params, AdamW for others. Arguments: param_groups: List of dicts with 'kind' = 'adamw' or 'mu

Sin/cos RoPE with an explicit ``positions`` argument so the same layer can handle both padded-batch training and NA inference.