---
title: "MARBLE: exact sequence parallelism for packed multimodal masks"
description: "An exact solution towards sequence parallelism for packed multimodal sequences: balanced ring attention with PTRR, an online FlashAttention-4 kernel, chunked ring attention and a sparse all-to-all, chosen per workload and layer."
canonical: "https://lumalabs.ai/news/packed-multimodal-attention"
source: "https://lumalabs.ai/news/packed-multimodal-attention.md"
---

# MARBLE: exact sequence parallelism for packed multimodal masks

Alex Redden, Akshat Tripathi, Leon Meng, Sam Pan, Michael Tu, Sourish Kundu, Narayanan Sundaram, Thomas Neff · Luma AI · October 19, 2026

An exact solution towards sequence parallelism for packed multimodal sequences: balanced ring attention with PTRR, an online FlashAttention-4 kernel, chunked ring attention and a sparse all-to-all, chosen per workload and layer.

## 00 · Introduction

In multimodal models, padding every sample to a fixed length is inefficient, so typical training frameworks pack many samples into one long sequence[1], letting the attention mask correctly handle cross-sample boundaries. Every sample within a sequence has its own attention mechanism: text is causal, images or video are usually bidirectional, generated images may be masked out from future samples and multiple images may be combined into a group[2]. Furthermore, any layer in the model may also use forms of sparse attention such as sliding window[3, 4] or video sparse attention[5]. Modern kernels such as FlashAttention-4 (FA4)[6] or FlexAttention[7] compose a hierarchy of block-sparse tiles with fine-grained masking functions.

When scaling to millions of tokens, we need to distribute sequences to fit them onto each rank's VRAM. Ulysses[8] distributes attention heads with all-to-all collectives, meaning that each rank will receive the whole sequence, requiring no change to the masking logic. However, its parallelism degree is limited by having to evenly divide the number of key/value heads. In contrast, ring attention[9] distributes the sequence by splitting it into key/value shards that are subsequently passed across ranks in multiple ring steps. When splitting heterogeneous sequences like this, the resulting shards can be unbalanced, resulting in stalls as slower ranks need to wait for faster ones. Additionally, merging partial results after each step in bf16 adds rounding error with every step. Unified sequence parallelism (USP)[10] composes the two approaches, enabling scaling beyond the Ulysses degree at the cost of inheriting drawbacks from both. Finally, FlashMask[11, 12] sidesteps the balancing issues of ring attention by all-gathering every key. While being efficient, at longer sequences this approach can run out of GPU memory due to storing the gathered keys. As a result, the optimal strategy changes within a model and even within a batch, depending on the layer and sample distribution.

*MARBLE chooses the communication and distribution scheme based on mask patterns for each layer. Top: example layouts for different multimodal recipes. Bottom: attention forward and backward speedup over FlashMask + Ulysses at each workload's largest measured scale.*

We introduce MARBLE (Mask-Aware Ring with Balanced, Lossless Execution), our approach for sequence parallelism that handles packed multimodal sequences via a *balanced* and *exact* ring attention core. The core of MARBLE is a block-sparse planning strategy that uses per-token metadata with processing-time round-robin (PTRR)[13] to balance the work evenly across ranks. Additionally, we built an online FA4 kernel that keeps the softmax state in fp32 across ring steps to avoid rounding errors, resulting in outputs and gradients matching single GPU results. Furthermore, we use a chunked ring attention variant for long generations, applying PTRR to contiguous chunks and running every step via dense varlen FA4. Our implementation can also decide to fetch, via a sparse all-to-all, only the key/value blocks any given rank's queries need, which for sliding-window attention is just the last window. The token order and communication primitives are chosen from the mask for each workload and layer.

When evaluated on five models across 1–8M tokens, running on 32–128 H100s, MARBLE outperforms a FlashMask + Ulysses hybrid in 47 of 48 cases (1.05–7.2×), outperforms FlashMask's context parallelism in all 43 cases it can run (1.22–48×), and outperforms MagiAttention in 17 of 18 cases. The speedup grows with the number of GPUs from a median of 1.28× at 32 GPUs to 2.63× at 128. When tested in end-to-end training, we reach step times of 13.2 s on Gemma 3 27B for 1M tokens on 32 GPUs against 32.5 s with FlashMask and 15.6 s with Ulysses alone. This post covers the main ideas; the full evaluation and more implementation details are in our upcoming technical report.

**Meet us at PyTorch Conference 2026.** We're in San Jose this week, and our team is hosting a happy hour on Tuesday, Oct 20, 7–10pm at Suspiro, near Santana Row. If you work on distributed training, GPU kernels, RL or inference, come talk shop: [request a spot here](https://luma.com/49muebiq).

## 01 · Packed multimodal masks

A packed sequence is a list of samples, where each sample itself contains several spans requiring their own attention mode. Within a sample, each span attends to itself according to its attention mode and later spans may attend to earlier ones, and across sample boundaries no sample attends another. We define five attention modes that cover all models we evaluate:

| Mode | Within span | Seen later | Example |
|---|---|---|---|
| causal | left to right | yes | text; Qwen3-VL image |
| full | bidirectional | yes | Gemma 3 crop; clean frame |
| noise | bidirectional | no | generated frame |
| group | all group spans | yes | Molmo2 images, frames |
| window | last W tokens | within W | Gemma 3 local layers |

Qwen3-VL[14] and Qwen3.8[15] read image and video tokens causally, like text. Gemma 3[4] uses a sliding window for five of every six layers, and uses bidirectional attention within each image span. Molmo2[2] combines all the image spans into one large bidirectional group, so earlier frames can attend to later frames and vice versa, but not to timestamps and text in between them. BAGEL[16] combines understanding and generation, where a generated image is a *noise* span that is ignored by all the other spans, followed by a clean copy that later frames attend to in cases of interleaved video generation. Together, the attention modes and sample boundaries define the *packed mask*: for every query/key pair, whether the query may attend the key.

*A packed multimodal mask. Rows are queries and columns are keys. Each legal attention pair within one span is colored by that span's attention mode, pairs where a later span reads an earlier one have their own color, and the strips give each token's attention mode.*

We encode the mask with four integers per token: a sample id, a segment id shared by the tokens of one bidirectional span, a noise-group id for generated spans, and the token's original position. This builds on top of the structure introduced by BAGEL[16], additionally adding the group and window modes, and the position track to ensure correctness after reordering tokens. The function below covers all attention modes, and it compiles into FA4 as a CuTe DSL `mask_mod`[6, 7]. For causal and sliding window modes, `W` is the sliding window, and `∞` for globally causal layers.

```python
causal(q, k)      = 0 <= pos[q] - pos[k] < W
same_span(q, k)   = seg[q] == seg[k] >= 0
keep_noise(q, k)  = noise[k] < 0 or noise[q] == noise[k]
same_sample(q, k) = sample[q] == sample[k]
mask_mod = and_masks(or_masks(causal, same_span),
                     keep_noise, same_sample)
```

## 02 · Two axes of parallelism

Ulysses[8] keeps the sequence split outside attention and switches from sharding on sequence to sharding on attention heads with an all-to-all before and after each attention call. Subsequently, every rank attends over the whole sequence for a subset of heads. As a result, it is perfectly load balanced across ranks, but requires that its degree evenly divides the head count. In contrast, ring attention[9] gives each of R ranks a chunk of queries and keys and rotates them over R steps. USP[10] composes the two, with ring degree R and Ulysses degree U: each rank holds 1/(RU) of the sequence outside attention, and Ulysses inside a node leaves each rank 1/R of the sequence for 1/U of the heads. Ring attention then rotates the Ulysses-sharded sequence R steps (denoted as USP R×U going forward).

*USP R×U: Ulysses splits attention heads across the ranks of a node, and ring attention splits the sequence across nodes and rotates the key/value shards one hop per step.*

## 03 · Exact, load-balanced ring attention

Ring attention is flexible enough to handle every attention type at every scale. Memory is bounded because each rank only holds a fixed number of key/value shards in memory at a time, and there is no dependence on the attention head count. Its main issues for packed multimodal masks are load balancing, inefficiencies due to communication, and accuracy when merging intermediate results. In this section, we present solutions for all these problems, and further carry them over to dense approaches. These solutions rely on a *planner* that runs once per packed batch, before the forward pass. It reads only the per-token metadata tracks and builds each rank's plan: the token order, and for every ring step the blocks or dense calls the rank computes. Every layer reuses the plan, forward and backward, and a data-loader worker can build it on the CPU while the previous step trains.

### Block classification

For each 128-row query block, FA4's block-sparse kernels receive a list of key/value blocks to operate on, each tagged *full* (no per-element masking) or *partial* (run the `mask_mod`). Similar to FlashMask[11], we classify block pairs from cheap per-block summaries to avoid materializing the mask. These summaries contain the minimum and maximum of each metadata track in a block, plus a small bitmap of the segments each block contains. A block pair is *full* if every pair is provably legal, *empty* if every pair is provably illegal, and *partial* otherwise.

*Block pairs are classified from per-block summaries of the metadata: full when every pair is provably legal, empty when every pair is provably illegal, and partial otherwise.*

Ring attention turns one self-attention mask into R submasks on each rank. At step s, rank r computes attention between its query shard and the key/value shard of rank (r − s) mod R. Because we reorder tokens to improve load balancing, a block's local index no longer corresponds to the position in the sequence. Therefore, each step requires separate query-side and key-side copies of the metadata tracks, including absolute positions. Instead of building a FlexAttention `BlockMask`[7], which evaluates the mask on every query/key pair, our planner classifies each step's block pairs from the per-block summaries alone, resulting in only one test per block pair instead of one per token pair. The implementation of the kernel does not need all four tracks: each key's visible queries form one range [lo<sub>k</sub>, hi<sub>k</sub>), so each ring step tests lo<sub>k</sub> ≤ pos<sub>q</sub> < hi<sub>k</sub> on partial tiles, which makes ring and sparse all-to-all steps 1.04–1.38× faster on 16 H100s at 2M tokens, with unchanged errors.

### PTRR: balancing ring attention steps

Ring attention needs to synchronize at every step, and as a result, a step costs as long as its slowest rank. Contiguous shards of a packed mask give the busiest rank 2–2.6× the average work on our synthetic packed masks. PTRR, a load balancer from PyTorch's context-parallel module[13], reorders query blocks before sharding so that ranks receive equal work. We sort blocks by cost, then assign them to ranks in rounds of R, reverse every other round, and keep each rank's blocks in sequence order for locality.

*PTRR on a 16-block example: sort the query blocks by cost, assign them to ranks in rounds of R with every other round reversed, and keep each rank's blocks in sequence order.*

Distributing blocks across ranks also balances out the work of each step, as every step pairs one rank's queries with another rank's keys. PTRR rebalances the busiest rank to within 0.1% and roughly halves the layer time.

*Each ring step lasts as long as its slowest rank. With contiguous shards, single steps are far off; after PTRR, every step does even work.*

**Zero-SM communication.** When naively using NCCL `send`/`recv`, each step's key/value transfer emits collective kernels that occupy the GPU's streaming multiprocessors (SMs), which subsequently cannot be fully utilized for the attention computation itself. We move these transfers to one-sided NCCL remote memory access (RMA) transfers via `nccl4py` `put_signal`/`wait_signal`. In this API, the receiver signals a free buffer, and the owner writes into it via the copy engine or the network interface card. Across two nodes this makes a 16-way ring 6% faster than NCCL `send`/`recv` (1433 vs 1524 ms); within one node, or when Ulysses carries most of the traffic, the two transports are within 1%.

*NCCL `send`/`recv` occupies SMs that attention needs. One-sided writes from the copy engine or the network card leave them to the kernel.*

### Exact ring attention accumulation

Standard ring attention runs a FlashAttention forward per step, outputting a normalized bf16 partial output and its log-sum-exp (LSE). Subsequently, it merges partial outputs with an LSE rescale. While this algebra is exact, it results in every step rounding partial outputs to bf16, and again rounding the running result when merging. For packed sequences, this error (against an fp32 reference) can be up to 56% higher compared to single GPU execution.

*MARBLE's online FA4 kernel. A per-step merge rounds bf16 partial outputs at every ring step. The online kernel carries the fp32 accumulator, row max m and row sum ℓ across steps and only casts once.*

As the FlashAttention forward already loops over key/value blocks with a running row max m, row sum ℓ and an unnormalized fp32 accumulator, each ring attention step is the same loop over a different set of blocks. The online FA4 kernel is built on top of FA4's Hopper forward. First, the fp32 accumulator, m and ℓ stay in high-bandwidth memory (HBM). Second, on every call but the first, the producer warp loads the previous fp32 output tile with a Tensor Memory Accelerator (TMA) load alongside the query tile, and the PV matmul accumulates onto it. Similarly, on every call but the first, we start from the stored m and ℓ instead of −∞ and 0. Finally, only the last call normalizes the output, writes the LSE and casts to bf16.

The backward pass can reuse the existing kernel. Each step runs FA4's backward against the final LSE, starting from an fp32 dQ accumulator that is carried across steps, while each shard's fp32 dK/dV accumulator travels from rank to rank. With the online kernel, ring attention's output and dQ error equal a single GPU's at every ring degree and on every layout we tested. One online call is about 4% slower than one FA4 call over the same keys; the main reason to use the online kernel is accuracy, not speed. On Blackwell, FA4's split-KV path already produces partial outputs in fp32 precision, and merging those in fp32 gives the same exactness without writing a new kernel.

*One rule for every mode: no partial result is rounded below the precision of a single attention call, and gradients summed across calls or ranks are summed in fp32.*

### Chunked ring attention

In long-video generation, where samples may exceed 128k–1M tokens of interleaved noise and clean frames, almost every tile is fully or partially attended, and dense variable-length (varlen) FA4 calls are faster than masked calls. MARBLE's chunked ring attention strategy cuts the sequence into contiguous chunks and runs every ring step as a few dense varlen calls. Against an earlier chunk, each sample's queries need one full-attention call over the chunk's non-noise keys, since no other span reads a noise span. Within the queries' own chunk, three kinds of calls cover the rest. Causal calls over the non-noise tokens cover text, as well as keys up to each query for bidirectional spans. The keys after each query come from doing a reverse causal call on the same span's tokens. Finally, generated image spans use one full attention call across their own tokens and clean tokens before it in the chunk.

*How the chunked ring attention covers a BAGEL generation mask with dense varlen calls. Each legal pair is colored by the call that computes it, and every pair is computed exactly once. A span that crosses a chunk boundary needs one more call at the ring step that brings the other chunk.*

Dense calls need contiguous chunks, so our chunked ring attention assigns a few large chunks per rank instead of PTRR's 128-token blocks. We choose the layout that minimizes the sum of each step's slowest rank. With W<sub>rj</sub> being the legal pairs between rank r's queries and rank j's keys, and rank r holding rank (r − s) mod R's keys at step s, a layout costs

```text
C = Σ_s  max_r  W[r, (r − s) mod R]  +  τ · R · k
```

where k is the number of contiguous ranges per rank and τ is the fixed cost of one range in one step. Each range is a separate varlen sequence in every step.

For the chunked ring attention, our planner tries 1, 2, 4, …, 64 chunks per rank to determine the best configuration. It chooses between two strategies to assign the chunks to ranks with PTRR's rule: once in the order of work, and once in the sequence order (which gives each rank an early and a late chunk, like a zigzag). Starting from the better of these two, it greedily swaps chunks between ranks whenever that lowers the per-step work, and finally keeps the chunk count with the lowest total cost C. The planner is run on every rank since it will deterministically produce the same result given the same random seed. We fit τ on nine BAGEL runs at 1–8M tokens, and C predicts the runtime eight out of nine times within 3%.

*Per-step balance on a toy BAGEL sequence. A pair's color is the ring step that computes it, and each step's slowest rank is outlined. The bars are the work of each step's slowest rank; their sum is the cost C without the overhead term.*

## 04 · Sparse all-to-all and strategy selection

Ring attention sends every key/value shard to every rank, whether or not that rank's queries can see it. When queries attend to only a few samples, the resulting ranks will not be utilized effectively, having to wait for other ranks. MARBLE therefore offers an alternative to ring attention, a sparse all-to-all, which also covers sliding-window layers. Both are exact and share the planner, metadata and kernels of ring attention.

*Ring attention, sparse all-to-all, and sparse all-to-all on a sliding-window layer, over one packed mask on four contiguous ranks. The colored strip under each mask marks the key blocks each primitive moves to rank 2.*

First, each rank computes exactly which remote 128-token key/value blocks it needs from its queries' key intervals, and fetches them with one variable-sized all-to-all. Then, it uses a single FA4 pass for both its local and fetched keys. Because each rank fetches only the blocks it needs, on our tested sequences we avoid sending 55–100% of the remote keys that ring attention would send. For the backward pass, dK and dV of the fetched blocks are returned with the inverse all-to-all and accumulated in fp32. For sliding-window layers, each rank needs at most W keys from the ranks before it. The sparse all-to-all then fetches exactly those, independent of the ring degree, which makes sliding-window layers such as in Gemma 3 cheap to parallelize. Within a node, Ulysses over NVLink is still the fastest approach in our tests, so we use USP R×U with U inside a node (4 or 8, as the heads allow) and R across nodes.

### Layout transforms

Each transform permutes tokens (since we can freely reorder them) once at the model input (activations, metadata and positions together). Every token carries its metadata and original position with it, and the mask function reads only those, so each query still attends exactly the same keys. PTRR scatters each sample over all ranks so that every ring step does even work. In contrast, contiguous chunks keep the sparse all-to-all communication volume small but lead to imbalanced ring attention. Therefore, we let ring attention run over PTRR blocks, and the sparse all-to-all over contiguous chunks.

For vision-group sequences like in Molmo2, we store all the text first and its group spans after in reverse, so that each key's visible queries form one range. By sorting like this, each key's visible queries become one contiguous range, and partially masked tiles on Molmo2 video at 512k tokens drop from 49% to 4%.

*Layout transforms permute tokens once at the model input. Ring attention runs over PTRR blocks, the sparse all-to-all over contiguous chunks, and vision-group order makes each key's visible queries one contiguous range.*

### Strategy selection

Which strategy we select follows from three properties of the mask and model: *locality* (what fraction of remote keys a rank's queries can see), *density* (whether visible tiles are full, partial or interleaved) and the *head budget* (how far we can extend the Ulysses degree).

| Workload | Mask property | Strategy |
|---|---|---|
| Packed understanding (documents, images, video) | strong locality | sparse all-to-all + contiguous chunks |
| Sliding-window layers | locality by construction | sparse all-to-all (last W keys) |
| Long dense generation (BAGEL video) | little locality, compute-bound | chunked ring attention, per-step balance |
| Hybrid sliding-window / global models | differs per layer | sparse all-to-all, planned per layer |
| More GPUs than key/value heads | any | USP, both degrees > 1 |

For a given number of GPUs and attention heads, the planner splits into a ring degree R and a Ulysses degree U (replicating key/value heads when needed by U). For each split, the planner distributes the sequence in two ways: ring attention or sparse all-to-all, each with its token order. The planner scores each candidate via the quantities it already computes: the compute time of the busiest rank, and the fitted throughput of the kernel it would use (block-sparse or dense). Furthermore, it considers the Ulysses all-to-all bytes and the ring attention or sparse all-to-all bytes, split into intra- and inter-node communication. Its eight hardware constants (kernel throughputs, bandwidth inside and across nodes, and fixed costs per ring step and per all-to-all) are fitted to our measured layer times by minimizing the squared relative error, so that short and long runs count equally for the strategy selection. Within ring attention, the planner switches to the chunked ring attention when a layout's median full or noise span is at least 512 tokens, as in long-video generation, and then picks its chunk layout by the cost C. For sliding-window layers, it uses the last-window all-to-all. Our top-ranked candidate is within 10% of the measured best in 93% of our 85 cases and within 5% in 86%.

### Planning cost

Every packed batch has a new layout and needs a new plan: the token order, each rank's block or call lists for every step, and all-to-all indices. A plan depends only on the layout and the rank's position in the ring and Ulysses groups, never on activations, and each rank builds its own plan to avoid communication across ranks. For training, a data-loader worker process can therefore build it asynchronously on the CPU while the previous step is in progress. At 4M tokens with USP 8×4 and 12 CPU threads, the slowest rank plans in 0.1–7.9 s depending on the mode, and 1.1 s for the chunked ring at 8M tokens. Training steps in our end-to-end runs take 6–44 s, meaning that planning can be overlapped with batch prefetching.

*A plan depends only on the layout, so a data-loader worker builds the next batch's plan on the CPU while the GPUs train on the current one.*

## 05 · Results

At scale we use H100 80GB nodes with 8 GPUs each (NVLink inside a node, InfiniBand across nodes), from 32 to 128 GPUs. Each number is the forward and backward time of one attention layer in bf16, unless marked end to end. We evaluate the attention of Qwen3-VL-30B-A3B, Qwen3.8-27B, Molmo2-8B, Gemma 3 27B and BAGEL-7B. Instead of synthetic masks, we build packed layouts of 1–8M tokens from each model's tokenizer and vision preprocessing rules, applied to document statistics of Molmo2's video data[2] and OBELICS[17] for understanding and BAGEL's token setup[16] for generation.

*Packed layouts built from each model's tokenizer and vision preprocessing rules.*

FlashMask CP is FlashMask's public PyTorch context parallelism[12], in which every rank gathers the whole key/value sequence by combining a FlashMask-4 kernel with a balanced all-gather. We also build FlashMask + Ulysses, which runs Ulysses inside each node and FlashMask CP across nodes, so each node only receives the key/value heads it needs. FlashMask-3's Hopper kernel[18] runs only through PaddlePaddle[19]: with PaddleFleet's balanced all-gather it is 1.1–2.2× slower than FlashMask CP at 32–64 GPUs. Additionally, when combined with Ulysses the same way, it is 1.1–1.9× slower than FlashMask + Ulysses. Its overlapped forward also does not close the gap, so the tables report the FlashMask-4 baselines. MagiAttention[20] balances mask slices across ranks with a dispatch solver, sends each key/value slice only to the ranks that need it (GroupCast) and sums the gradients back onto their owners (GroupReduce). On single nodes we also compare against TransformerEngine[21], yunchang[22] and PyTorch's context parallelism[23], which support causal and document masks only. We evaluate on every system's fastest configuration after doing a sweep: ours over ring and Ulysses degrees, communication primitives and layouts in exact mode, FlashMask CP over its load-balancing options, FlashMask + Ulysses over Ulysses degrees, and each library over all its modes. For each implementation we also perform a forward and backward check. In these checks, we distribute an 8k-token sequence with its own layout, and we compare output, dQ, dK and dV on every rank's tokens against an fp32 dense reference. 51 of 52 baseline configurations pass; yunchang 0.6.4 fails on causal masks (0.57 relative error) because its FlashAttention-3 wrappers drop the `causal` flag, and we benchmark it with a one-line patch.

### Accuracy

We compare each attention method and mode (outputs and gradients) against an fp32 call over the whole sequence. We test every mask type on 32k-token packs with 4 key/value heads and USP 4×2 on 8 H100s. Every mode we use at scale is within 2% of a single GPU. FlashMask CP and MagiAttention's default bf16 reduction add 11–22% dK/dV error. On 8 GPUs with a causal mask and 8k tokens, the ring attention of TransformerEngine, yunchang and PyTorch's context parallelism has 21–55% more output error and 64–114% more dQ error than a single GPU, while ours matches it. A 200M-parameter language model trained for 4000 steps with each ring variant ends no more than 0.0001 above the output loss of a Ulysses run.

### Scaling to 128 GPUs

We compare our fastest exact configuration against each baseline's fastest across 48 (model, layout, scale) cases with a FlashMask + Ulysses result. We are faster in 47 and equal in one, and FlashMask CP is slower than us in all 43 cases it can run. The relative speedup grows with scale, from a median of 1.28× at 32 GPUs to 2.63× at 128 GPUs. This can be attributed to the baselines' communication volume growing with the sequence length, while our implementation grows with the number of valid elements in the mask. At 4M tokens on 64 GPUs, FlashMask CP runs out of memory on Gemma 3 and Molmo2, where our implementation uses 7–9 GiB. MagiAttention, even when tuning its chunk size and overlap degree, is 1.05–12.3× slower than MARBLE in 17 cases and faster in one, BAGEL's packed training mix at 2M tokens (0.90×). The equal case is Qwen3.8 video at 1M tokens on 32 GPUs (0.99×): its head dimension of 256 has no block-sparse FA4 kernel on H100, so we run dense calls.

*The margin grows with scale. Speedup of MARBLE over FlashMask + Ulysses; lines join per-model medians, faint dots are individual layouts. Dashed: FlashMask CP's speed relative to FlashMask + Ulysses (median over cases it can run).*

### Long-video generation

BAGEL long video is the workload where the sparse all-to-all loses: samples of 128k–1M tokens with little locality, bound by compute. The chunked ring attention with per-step balance is the fastest exact mode at every length. Refining the chunk layout is worth 7% at 8M tokens and 4% at 4M against the plain 64-chunk assignment.

| Tokens / GPUs | Ours (ms) | GiB | FlashMask + Ulysses | FlashMask CP | Speedup |
|---|---|---|---|---|---|
| 1M / 32 | **407** | 4.8 | 449 | 497 | 1.10× |
| 2M / 32 | **1015** | 9.0 | 1148 | 1241 | 1.13× |
| 4M / 32 | **1861** | 19.3 | 2102 | 2418 | 1.13× |
| 8M / 64 | **2190** | 18.5 | 2412 | 2856 | 1.10× |

*BAGEL long video: chunked ring attention with refined per-step balance against FlashMask (forward + backward of one layer, ms; peak memory per GPU).*

### End to end

We time full forward and backward training steps of Gemma 3 27B, Molmo2-8B and BAGEL-7B on real packs, with attention planned per model by the chosen strategy. On Gemma 3 at 2M tokens, one ring plan for all layers takes 17.1–17.7 s per step; per-layer plans with a last-window all-to-all for the local layers bring this to 15.0–15.3 s, and a sparse all-to-all for the global layers to 13.3–13.7 s. BAGEL-7B on long-video packs runs the chunked ring at 64k tokens per GPU, 1.13× and 1.10× faster end to end than FlashMask + Ulysses, the same margins as attention alone.

| Model | Tokens / GPUs | Ours | FlashMask CP | Ulysses |
|---|---|---|---|---|
| Gemma 3 27B | 1M / 32 | **13.2 s** | 32.5 s (60 GiB) | 15.6 s |
| Gemma 3 27B | 2M / 64 | **13.3 s** | out of memory | n/a |
| Molmo2-8B | 1M / 32 | **6.05 s** | 11.5 s (30 GiB) | 7.65 s |
| Molmo2-8B | 2M / 64 | **7.39 s** | 20.1 s (45 GiB) | n/a |
| Molmo2-8B | 4M / 64 | **14.0 s** | out of memory | n/a |
| BAGEL-7B | 2M / 32 | **43.5 s** | 51.9 s (39 GiB) | n/a |
| BAGEL-7B | 4M / 64 | **43.6 s** | 53.8 s (54 GiB) | n/a |

*End-to-end training step time on real packs, excluding planning. The FlashMask CP column shows peak memory.*

On one node of 8 B300s, the same configurations run 1.7–2.3× faster, on newer kernels, and Blackwell needs no new kernel for exactness. The upcoming technical report has the full results, including single-node comparisons with TransformerEngine, yunchang and PyTorch's context parallelism, the ablations of each piece, and Blackwell.

## 06 · Limitations and conclusion

Planning takes 0.1–8 s per batch on the CPU, against 0.1–1.0 s on the GPU for FlashMask; it hides behind the previous step only if the data loader has CPU time to spare. With full multi-head attention, summed dK and dV stay at 1.3–1.4× a single GPU's error; none of the evaluated models uses it. The online kernel is Hopper-only (Blackwell gets exactness from fp32 split-KV partials). FlashMask + Ulysses is our own construction, and FlashMask's authors may build such a hybrid differently. Our scale comparisons use real packs built from tokenizer rules and dataset statistics rather than recorded training batches.

Packed multimodal masks have no single best parallel strategy. Locality favors fetching only what the mask can see, sliding windows favor fetching just the window, long dense generation favors ring attention, and the head count decides how far Ulysses can go inside a node. MARBLE covers all of them from one description of the mask, one planner and one set of kernels, and picks among them per workload and layer. Ring attention at its center makes this work: PTRR balances it, its per-step objective carried to dense chunks wins long-video generation, and the online kernel makes it exact, with the same rule keeping every other mode exact. The result is faster than FlashMask, a FlashMask + Ulysses hybrid and MagiAttention in nearly every case we measured, with a margin that grows with scale, while computing the same attention as a single GPU. Our code, layouts and benchmarks will be released.

If you're at PyTorch Conference this week, come find us at our happy hour on Tuesday, Oct 20, 7–10pm at Suspiro in San Jose: [request a spot](https://luma.com/49muebiq).

## 07 · References

1. M. M. Krell, M. Kosec, S. P. Perez, A. Fitzgibbon, "Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance," 2021. [arXiv:2107.02027](https://arxiv.org/abs/2107.02027)
2. C. Clark et al., "Molmo2: Open Weights and Data for Vision-Language Models with Video Understanding and Grounding," 2026. [arXiv:2601.10611](https://arxiv.org/abs/2601.10611)
3. I. Beltagy, M. E. Peters, A. Cohan, "Longformer: The Long-Document Transformer," 2020. [arXiv:2004.05150](https://arxiv.org/abs/2004.05150)
4. Gemma Team, "Gemma 3 Technical Report," 2025. [arXiv:2503.19786](https://arxiv.org/abs/2503.19786)
5. P. Zhang et al., "VSA: Faster Video Diffusion with Trainable Sparse Attention," 2025. [arXiv:2505.13389](https://arxiv.org/abs/2505.13389)
6. T. Zadouri, M. Hoehnerbach, J. Shah, T. Liu, V. Thakkar, T. Dao, "FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling," 2026. [arXiv:2603.05451](https://arxiv.org/abs/2603.05451)
7. PyTorch Team, "FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention," 2024. [pytorch.org/blog/flexattention](https://pytorch.org/blog/flexattention/)
8. S. A. Jacobs et al., "DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models," 2023. [arXiv:2309.14509](https://arxiv.org/abs/2309.14509)
9. H. Liu, M. Zaharia, P. Abbeel, "Ring Attention with Blockwise Transformers for Near-Infinite Context," 2023. [arXiv:2310.01889](https://arxiv.org/abs/2310.01889)
10. J. Fang, S. Zhao, "USP: A Unified Sequence Parallelism Approach for Long Context Generative AI," 2024. [arXiv:2405.07719](https://arxiv.org/abs/2405.07719)
11. G. Wang et al., "FlashMask: Efficient and Rich Mask Extension of FlashAttention," 2025. [arXiv:2410.01359](https://arxiv.org/abs/2410.01359)
12. PaddlePaddle, "FlashMask," GitHub repository (FlashMask-4 CuTe DSL kernels, PyTorch interface, context-parallel load balancing). [github.com/PaddlePaddle/FlashMask](https://github.com/PaddlePaddle/FlashMask)
13. PyTorch Team, "Context Parallel _PTRRLoadBalancer (torch.distributed.tensor.experimental)," 2025. [github.com/pytorch/pytorch](https://github.com/pytorch/pytorch/blob/161d490730c9646054a4c66f7713f35890e55353/torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py#L295-L467)
14. Qwen Team, "Qwen3-VL Technical Report," 2025. [arXiv:2511.21631](https://arxiv.org/abs/2511.21631)
15. Qwen Team, "Qwen3.8-27B," model card, 2026. [huggingface.co/Qwen/Qwen3.8-27B](https://huggingface.co/Qwen/Qwen3.8-27B)
16. C. Deng et al., "Emerging Properties in Unified Multimodal Pretraining," 2025. [arXiv:2505.14683](https://arxiv.org/abs/2505.14683)
17. H. Laurençon et al., "OBELICS: An Open Web-Scale Filtered Dataset of Interleaved Image-Text Documents," 2023. [arXiv:2306.16527](https://arxiv.org/abs/2306.16527)
18. G. Wang et al., "FlashMask-3: Efficient and Expressive Mask-Aware Distributed Attention," 2026. [arXiv:2610.17901](https://arxiv.org/abs/2610.17901)
19. Y. Ma, D. Yu, T. Wu, H. Wang, "PaddlePaddle: An Open-Source Deep Learning Platform from Industrial Practice," Frontiers of Data and Computing, 2019.
20. Z. Tao, Y. Huang, "MagiAttention: A Distributed Attention Towards Linear Scalability for Ultra-Long Context, Heterogeneous Mask Training," 2025. [github.com/SandAI-org/MagiAttention](https://github.com/SandAI-org/MagiAttention)
21. NVIDIA, "Transformer Engine," `DotProductAttention` context parallelism. [github.com/NVIDIA/TransformerEngine](https://github.com/NVIDIA/TransformerEngine)
22. J. Fang, S. Zhao, "YunChang: long-context-attention," reference implementation of USP. [github.com/feifeibear/long-context-attention](https://github.com/feifeibear/long-context-attention)
23. X. Wu, C.-C. Huang, "Introduction to Context Parallel," PyTorch Tutorials, 2025. [docs.pytorch.org/tutorials/unstable/context_parallel.html](https://docs.pytorch.org/tutorials/unstable/context_parallel.html)

## Citation

If you use this work, please cite it as:

> Alex Redden, Akshat Tripathi, Leon Meng, Sam Pan, Michael Tu, Sourish Kundu, Narayanan Sundaram, Thomas Neff. "MARBLE: exact sequence parallelism for packed multimodal masks." Luma AI, October 2026.

```bibtex
@misc{redden2026marble,
  title        = {{MARBLE}: Exact Sequence Parallelism for Packed Multimodal Masks},
  author       = {Redden, Alex and Tripathi, Akshat and Meng, Leon and Pan, Sam and Tu, Michael and Kundu, Sourish and Sundaram, Narayanan and Neff, Thomas},
  year         = {2026},
  month        = oct,
  howpublished = {Luma AI blog},
}
```