Encode per-sample ``input_ids`` into per-sample latents. Realizes the inference encoder ``q_phi(z_0 | x)`` of Eq. 2.1.1 of the paper.

Decode flattened latents into vocabulary logits. Realizes the conditional decoder ``p_theta(x | z_0)`` of Eq. 2.1.1. During block-wis

Set attention backend: 'naive' (default), 'sdpa', or 'flex'.

Accept ``(B, 1)`` or ``(B,)`` per-sample length tensors uniformly.

Additive block-diagonal block-causal attention mask for NA sequences. Invariants the inference code guarantees and that we rely on here: * `

Attention mask for the 2L training trick. Each sample has layout ``[clean(L) | noisy(L)]``. Q may cover the full ``2L`` (when ``txt_q_shape

FlexAttention block mask for the 2L training trick (NA form). Uses Fast-dLLM v2 layout ``[xt(noisy, L) | x0(clean, L)]`` per sample. The call

BOS-aligned best-fit packing dataloader for pretraining. Args: tokenizer_path: path to tokenizer.json (OLMo 2 / HuggingFace tokenize

Build a ``(B, 1)`` shape tensor from a Python list of per-sample lengths.

End-to-end Cola DLM inference (Eq. 2.2.4–2.2.6 of the paper). Realizes the three-step inference algorithm of *Continuous Latent Diffusion Lan

Run forward pass with specified layout and extract noisy predictions.

``List[Tensor(*_, c)]`` → ``(Tensor(L_total, c), txt_shape (B, 1))``.

Inverse of :func:`_flatten`: return a Python list of per-sample tensors.

Check if the VAE can reconstruct special tokens in context.

One FM training step on a batch. batch: list of (z_0, roles_str, L, sample_bs) tuples from prepare_sample. Returns: scalar loss.

Concat per-sample RoPE frequencies along the flattened sequence. ``txt_shape`` is ``(B, 1)`` and ``offset`` is a Python list of per-s

Build RoPE frequencies from explicit integer position IDs. ``position_ids`` is a 1-D ``(L,)`` long tensor of per-token positions — e.

Convert (B, T) token IDs to list of (z_0, L, block_size) tuples.

``txt_shape: (B, 1)`` → ``positions: (1, L_total)`` where each sample's positions restart from 0.

Per-sample Q positions aligned to the TAIL of K within each sample.

Infinite iterator over document text batches from parquet files.

Run embedding + patch ``Conv1d`` one sample at a time. ``nn.Conv1d`` requires a fixed length within a batch; looping over the (small)

Sinusoidal timestep embedding. Matches diffusers ``get_timestep_embedding`` with ``flip_sin_to_cos=False, downscale_freq_shift=0`` — i.e. the

Apply pre-computed ``freqs_q / freqs_k`` to ``txt_q / txt_k`` of shape ``(L, h, d)``. Not named ``apply`` because that would shadow :

Construct the noisy copy with boundary-aware noising. Returns: z_noisy: (L, d) loss_mask: (L,) target: (L, d) velocity ta

All-token noisy construction with optional simulated [P,R] boundaries. With probability prompt_block_prob per block, a random split creates a

Reconstruction loss: encode → decode → cross-entropy.

Cumulative sequence lengths with a leading zero, ``int32``.

Generate text block-by-block, yielding decoded text per block. Args: prompt_ids: tokenized prompt (list of int token IDs)

Return Cola-DLM cache directory, default to ./cache/.