DreamDojo Documentation

Multi-Dimensional Attention requires a specific tensor layout. In addition to the standard contiguous heads-last layout, the "sequence length" dimension must be unrolled into its original spatial/temporal representation.

Important Requirements:

1. Token Layout Shape: The dimensions representing the token layout (e.g., $X, Y$ for images) must be present in the tensor shape.
2. Shape Matching: The shapes of query, key, and value must match exactly along all dimensions of the token layout shape. This is because the API assumes query and context coordinate spaces are identical.

Supported Layouts

• 1-D (Language/Audio): (batch, X, heads, head_dim)
• 2-D (Images): (batch, X, Y, heads, head_dim)
• 3-D (Videos/3-D Images): (batch, X, Y, Z, heads, head_dim)

# 1-D case: language, audio
batch, X, heads, head_dim = query_1d.shape

# 2-D case: images
batch, X, Y, heads, head_dim = query_2d.shape

# 3-D case: videos / 3-D images
batch, X, Y, Z, heads, head_dim = query_3d.shape

# Requirement: query, key, and value must match on layout dimensions
assert query_2d.shape[1:3] == key_2d.shape[1:3] == value_2d.shape[1:3]

Imaginaire Attention requires input tensors (query, key, and value) to follow a specific memory layout: heads-last torch contiguous (torch.contiguous_format).

Tensor Shape Requirements: Inputs must be rank-4 tensors with the following dimensions:

1. Batch size
2. Sequence length
3. Number of attention heads
4. Head dimension

Memory Layout Details: This layout is consistent with PyTorch's contiguous_format, where the right-most dimension (head dimension) is the major dimension (stride 1). Tokens from different heads are interleaved in memory.

To ensure your tensors are correctly formatted, you can use the following verification logic:

def verify_heads_last_contig_tensor(x: Tensor):
    assert x.shape[0] == batch
    assert x.shape[1] == seqlen
    assert x.shape[2] == heads
    assert x.shape[3] == head_dim

    assert x.stride(3) == 1
    assert x.stride(2) == head_dim
    assert x.stride(1) == heads * head_dim
    assert x.stride(0) == heads * head_dim * seqlen

Use uv to manage the environment. Note that Blackwell architecture users must use Docker; virtual environment support for Blackwell is currently in development.

1. Install Git LFS and clone the repository:

   sudo apt install git-lfs
   git lfs install
   git clone git@github.com:nvidia-cosmos/<repository_name>.git
   cd <repository_name>
   git lfs pull

2. Install system dependencies:

   sudo apt install curl ffmpeg tree wget

3. Install uv:

   curl -LsSf https://astral.sh/uv/install.sh | sh
   source $HOME/.local/bin/env

4. Sync the environment with the desired CUDA extra:
   • For CUDA 12.8: --extra=cu128
   • For CUDA 13.0: --extra=cu130

To create a new environment:

uv sync --extra=cu128
source .venv/bin/activate

To install into an existing active environment (like Conda):

uv sync --extra=cu128 --active --inexact

Imaginaire Attention supports variable length sequences via two methods.

Option 1: Direct sequence lengths (Less efficient)

Pass seqlens_Q and seqlens_KV directly as tensors. This method is less efficient because it manually computes maximum sequence lengths and cumulative sums (including additional padding) during every call.

Option 2: Pre-computed parameters (More efficient)

For better performance, compute the cumulative sequence lengths and maximums once using generate_varlen_parameters and pass these pre-computed values to the attention layer. This avoids redundant computations in subsequent layers.

from cosmos_predict2._src.imaginaire.attention.varlen import generate_varlen_parameters

# Pre-compute parameters once
(
    cumulative_seqlen_Q,
    cumulative_seqlen_KV,
    max_seqlen_Q,
    max_seqlen_KV,
) = generate_varlen_parameters(query, key, value, seqlens_Q, seqlens_KV)

# Use pre-computed parameters in attention layers
output = attention(
    query=query,
    key=key,
    value=value,
    cumulative_seqlen_Q=cumulative_seqlen_Q,
    cumulative_seqlen_KV=cumulative_seqlen_KV,
    max_seqlen_Q=max_seqlen_Q,
    max_seqlen_KV=max_seqlen_KV,
)