After spending a week on image diffusion inference, the natural next step was video. And the moment I tried to run a video diffusion model end to end, I hit the wall that everyone hits: the token count explodes. A single 720p frame patchified at 2x2 gives you roughly 900 spatial tokens. Multiply by 16 frames at 8 fps and you are looking at 14,400 tokens in a single forward pass. At 48 frames, that is 43,200 tokens. The self-attention alone would require a matrix of 43,200 x 43,200, nearly 1.9 billion entries per head. No single GPU can handle that in FP16 without running out of memory.
This is where context parallelism (CP) comes in, and it turns out to be one of the most elegant ideas in modern distributed inference.
Why tensor parallelism is not enough
Tensor parallelism (TP) splits the model weights across GPUs. Each GPU holds a shard of every linear layer and they collectively compute the full output through AllReduce. This works beautifully for LLMs because the sequence length is manageable and the weights are the big cost.
Video models flip the equation. The weights of a typical video DiT (Diffusion Transformer) might be 5 to 10 billion parameters, fitting comfortably on one or two GPUs. But the activations scale with sequence length squared in attention, and sequence length scales with frames times spatial resolution. TP does not help here because it does not reduce the per-GPU activation memory. You need to split the sequence itself.
Context parallelism: splitting the sequence
Context parallelism partitions the token sequence across GPUs along the sequence dimension. If you have 43,200 tokens and 4 GPUs, each GPU holds roughly 10,800 tokens. Each device computes Q, K, V projections for its local tokens, then participates in a distributed attention step.
The key algorithm that makes this work is Ring Attention. The idea is simple and clever:
- Each GPU computes attention for its local Q against its local K, V.
- Then K, V blocks are passed around the ring to the next GPU.
- Each GPU accumulates partial attention outputs using the online softmax trick (the same trick that powers FlashAttention).
- After P steps (where P is the number of GPUs), every GPU has seen every K, V block and holds the correct output for its local Q tokens.
The total communication is P-1 rounds of sending K, V tensors of size (local_seq_len, head_dim) per head. Because the sends overlap with compute on the current block, you can hide most of the communication latency behind useful work.
The spatial-temporal twist
Video tokens have two-dimensional structure: spatial position within a frame and temporal position across frames. Smart implementations exploit this by partitioning along the temporal axis first. This means each GPU gets a contiguous chunk of frames rather than a scattered set of patches.
Why does this matter? Because many video architectures use factored attention: spatial attention within each frame and temporal attention across frames. If you split temporally:
- Spatial attention is fully local to each GPU (each GPU has complete frames).
- Only temporal attention requires cross-GPU communication.
This cuts the communication cost dramatically compared to a naive sequence split. NVIDIA's implementation in their Cosmos video generation pipeline uses exactly this strategy.
What the numbers look like
Consider generating a 5-second clip at 720p, 24 fps with a DiT that uses 2x2 spatial patches and a temporal patch size of 1. That gives us:
# Token count calculation
spatial_h = 720 // 2 # 360 spatial patches vertically
spatial_w = 1280 // 2 # 640 spatial patches horizontally
# With further downsampling in the VAE (8x), effective spatial tokens:
spatial_tokens = (720 // 16) * (1280 // 16) # 45 * 80 = 3,600
temporal_tokens = 24 * 5 # 120 frames
total_tokens = spatial_tokens * temporal_tokens # 432,000
# Attention memory per head (FP16, no FlashAttention):
attn_bytes = total_tokens * total_tokens * 2 # ~373 GB per head
# Obviously impossible on a single GPUWith FlashAttention, you eliminate the materialized attention matrix, but you still need to store all the K, V, and intermediate activations. At 432,000 tokens with hidden dimension 4096 and 32 heads, just the Q, K, V tensors take roughly 32 GB in FP16. Add the FFN activations and you are well past single-GPU territory.
With 8-way context parallelism, each GPU handles 54,000 tokens. The local K, V fit comfortably, and Ring Attention streams the rest.
Combining CP with other parallelism strategies
In practice, video generation uses multiple parallelism strategies simultaneously:
- Context parallelism (CP) across the sequence dimension to handle token count.
- Tensor parallelism (TP) across the hidden dimension if the model is large enough.
- Pipeline parallelism (PP) across DiT blocks if the model is very deep.
A typical Cosmos-1.0 setup uses CP=8 across 8 GPUs within a node (connected by NVLink for fast ring communication) and TP=1 because the model weights fit on each GPU. If scaling to multiple nodes, you might add TP=2 across nodes, giving you 16 GPUs total with CP=8, TP=2.
The rule of thumb: use CP first for video models because the bottleneck is sequence length, not model size. Add TP only when the model parameters themselves do not fit.
The communication pattern
Ring Attention communicates K, V blocks in a ring topology. For CP=8 with local sequence length L per GPU and head dimension d=128:
# Per-step communication (one direction of the ring)
kv_bytes_per_step = 2 * L * d * 2 # K and V, FP16
# With L = 54,000 and d = 128:
kv_bytes_per_step = 2 * 54000 * 128 * 2 # ~27.6 MB per head
# Total across 32 heads:
total_per_step = 27.6e6 * 32 # ~884 MB per ring step
# 7 ring steps for CP=8:
total_comm = 884e6 * 7 # ~6.2 GB totalOn NVLink with 450 GB/s bandwidth, that 6.2 GB takes about 14 ms. If each ring step's compute takes 20+ ms, the communication is fully hidden. This is why intra-node NVLink is critical for context parallelism.
Where this is heading
The trend in video generation is clear: longer sequences, higher resolution, more frames. Sora, Veo, Cosmos, and Kling all push toward minute-long, high-resolution video. Context parallelism is not optional for these workloads; it is the foundational distributed strategy.
The interesting frontier is combining CP with diffusion-specific optimizations like classifier-free guidance parallelism (running the conditional and unconditional passes on separate GPU groups) and temporal sliding windows where you generate overlapping chunks of frames and blend them.
Tomorrow I will look at multi-modal batching, where we tackle the challenge of batching requests that mix text, images, audio, and video in the same serving system.