Markdown

Variable-length attention and data packing¶

Scope: how to eliminate padding-token waste when training transformers on variable-length data, using two techniques that compose: variable-length (varlen) FlashAttention, which runs attention only over valid tokens instead of a padded rectangle, and data packing, which concatenates several short samples into one near-maximum-length sequence so there is almost nothing to pad. Both matter most for Vision-Language-Action (VLA) models, whose image-patch and instruction lengths vary sample to sample. Distilled from the JD Technology embodied-infrastructure report (Guo et al., 2026), which combines them for a 1.88x training throughput gain, plus that report's π0.5 visual-token-pruning case study. Companion to FlashAttention and MLA and the PyTorch attention APIs.

The numpy block labelled core math is self-contained, runnable, and validated; it reproduces the report's measured savings curve (2.28% to 89.73% as padding rises) and proves the block-diagonal mask that makes packing correct. The torch snippets are reference templates: pin versions and validate against your model before production. The report is one industrial source; its multipliers are that team's measurements, not portable guarantees.

What it is¶

Transformers process a batch as a rectangular tensor of shape [batch, seq_len, hidden], but real samples have different lengths. The default fix is padding: pick a fixed length L, and fill every shorter sample up to L with dummy tokens. Those padding tokens are masked out of the loss, but in a naive attention implementation they still occupy rows and columns of the L x L attention matrix, so the GPU spends real FLOPs and memory bandwidth on tokens that contribute nothing.

Two techniques remove that waste:

Variable-length FlashAttention (the flash_attn_varlen interface) takes the batch as one concatenated sequence plus a cu_seqlens array of cumulative sequence boundaries, and computes attention block by block only within each real sequence. No padded rectangle is ever materialized. For VLA models this is applied on both sides of the backbone: the vision encoder (for example Qwen2.5-VL's ViT) calls the varlen interface over valid visual tokens, and the language decoder concatenates visual features with text embeddings and runs varlen attention over the mixed sequence.
Data packing attacks the same waste earlier, during preprocessing: instead of one sample per sequence padded to L, it concatenates multiple short samples end to end into a single sequence close to L, so the trainer receives a no-padding (or near-zero-padding) stream. Packing needs varlen attention (or an equivalent block-diagonal mask) to keep the packed samples from attending across each other's boundaries.

They are complementary: packing minimizes how much padding exists; varlen attention makes whatever structure remains cheap to compute.

Why use it¶

Padding is pure waste on variable-length data. With fixed-length attention, the fraction of compute spent on padding roughly equals the fraction of positions that are padding. The report measured varlen time savings rising from 2.28% at a 3% padding rate to 89.73% at a 90% padding rate: the savings track the padding fraction almost exactly (validated below).
The win grows with sequence length. Attention is quadratic, so at long sequences varlen attention saves more than the linear padding fraction, and its achieved TFLOPS meets or exceeds the fixed-length path. The report saw varlen latency 25% to 90% lower than fixed-length for sequences beyond 8k.
Packing turns many short samples into few full ones. Multimodal instruction data is full of short samples; packing them into full sequences cut the number of sequences and drove padding toward zero, yielding a 1.88x throughput gain and a 46.87% total training-time reduction, with downstream accuracy flat or slightly up.
Accuracy is preserved. Because masking (or varlen boundaries) makes packed/varlen computation mathematically identical to the padded version for valid tokens, benchmark scores held steady across the report's tests.

When to use it (and when not)¶

Use varlen attention whenever your batch has non-trivial length variance and you are already on a FlashAttention-capable path. It is close to free and correctness-preserving. It is the default for multimodal VLA training.
Use data packing when you have many samples much shorter than the context window (instruction tuning, multimodal QA, short-trajectory embodied data). Packing's benefit is proportional to how much padding it removes.
Skip packing when samples already fill the context. If every sample is near L, there is nothing to pack, and packing only adds bookkeeping.
Be careful with packing when cross-sample leakage would corrupt training. Packing is only correct with a block-diagonal attention mask (or varlen boundaries) and correct position-id resets per sample. Without them, sample A attends to sample B and the loss is wrong. This is the load-bearing detail validated below.
Watch loss-normalization semantics. Packing changes how many samples share a sequence, so per-token vs per-sample loss averaging can shift; verify your reduction is what you intend.

Architecture¶

Both techniques share one idea: never spend compute on a token that is not real, and never let a real token attend to a token from a different sample. Varlen attention enforces this with cu_seqlens boundaries; packing enforces it with a block-diagonal mask over the concatenated sequence.

flowchart TB
  subgraph FIX["Fixed-length padding (wasteful)"]
    S1["sample lengths 300, 512, 180"] --> PAD["pad all to L=512"]
    PAD --> RECT["dense L x L attention per sample: padding computed"]
  end
  subgraph VAR["Variable-length + packing (ours)"]
    S2["same samples"] --> PACK["pack into one sequence near L"]
    PACK --> CU["cu_seqlens boundaries + block-diagonal mask"]
    CU --> BLK["attention only within each sample: no padding, no leakage"]
  end

Core math (runnable): savings curve, packing, and the leakage mask¶

This numpy block reproduces the report's central claims and the one correctness invariant. It shows (A) varlen savings track the padding-token fraction at moderate lengths but exceed it once attention dominates; (B) packing drives padding from over 50% to under 5% and the 1.88x throughput equals a 46.87% time reduction; (C) the block-diagonal mask zeroes all cross-sample attention while a naive full-attention softmax leaks across boundaries; and the π0.5 pruning arithmetic. Run: python3 varlen_packing.py.

# varlen_packing.py -- variable-length attention + data packing, core math. numpy only.
import numpy as np

# ---------- (A) Padding-compute model: linear (projections/MLP) + quadratic (attention) ----------
# Per sequence of length n: cost = alpha*n (token-linear) + beta*n^2 (attention, quadratic).
def batch_cost(lengths, L, alpha, beta):
    lengths = np.asarray(lengths, float)
    fix = len(lengths) * (alpha * L + beta * L * L)             # every seq padded to L
    var = np.sum(alpha * lengths + beta * lengths * lengths)    # compute only valid tokens
    return fix, var

def padding_rate(lengths, L):
    lengths = np.asarray(lengths, float)
    return 1.0 - lengths.sum() / (len(lengths) * L)             # fraction of positions that are padding

rng = np.random.default_rng(0)
L = 4096
lo = rng.integers(3900, 4096, size=256)                        # low padding (~3%)
hi = rng.integers(300, 520, size=256)                          # high padding (~90%)
p_lo, p_hi = padding_rate(lo, L), padding_rate(hi, L)
assert abs(p_lo - 0.03) < 0.02 and abs(p_hi - 0.90) < 0.02

# Token-linear-dominated regime (alpha >> beta*L): varlen savings ~= padding-token fraction.
a, b = 1.0, 1e-6
for lengths, p in [(lo, p_lo), (hi, p_hi)]:
    fix, var = batch_cost(lengths, L, a, b)
    assert abs((1 - var / fix) - p) < 0.02                     # matches paper: 3%->~2.3%, 90%->~89.7%
s_lo = 1 - batch_cost(lo, L, a, b)[1] / batch_cost(lo, L, a, b)[0]
s_hi = 1 - batch_cost(hi, L, a, b)[1] / batch_cost(hi, L, a, b)[0]
assert s_hi > s_lo                                             # savings monotonic in padding rate

# Attention-dominated regime (beta*L >> alpha): quadratic term makes savings EXCEED the token fraction.
saving_quad = 1 - batch_cost(hi, L, 1.0, 1.0)[1] / batch_cost(hi, L, 1.0, 1.0)[0]
assert saving_quad > p_hi

# ---------- (B) Data packing: first-fit-decreasing into bins of size L ----------
def pack(lengths, L):
    bins = []
    for n in sorted(lengths, reverse=True):
        assert n <= L, "a sample longer than the context cannot be packed; split or truncate it"
        for i, used in enumerate(bins):
            if used + n <= L:
                bins[i] += n; break
        else:
            bins.append(n)
    return bins

samples = rng.integers(50, 800, size=2000)
tokens = int(samples.sum())
bins = pack(samples, L)
pad_before = 1 - tokens / (len(samples) * L)                   # one sample per padded sequence
pad_after = 1 - tokens / (len(bins) * L)                       # packed sequences
assert len(bins) < len(samples) / 4                            # far fewer sequences
assert pad_before > 0.5 and pad_after < 0.05                   # padding driven from >50% to <5%

raised = False                                                 # edge: over-length sample must be refused
try: pack([L + 1], L)
except AssertionError: raised = True
assert raised

tput = 1.88                                                    # paper's "188%" is a 1.88x throughput
assert abs((1 - 1 / tput) - 0.4687) < 0.001                    # => 46.87% time reduction, as reported

# ---------- (C) Block-diagonal mask makes packing correct ----------
def softmax(x, axis=-1):
    x = x - x.max(axis=axis, keepdims=True)
    e = np.exp(x); return e / e.sum(axis=axis, keepdims=True)

seg = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2])                    # three packed samples in one sequence
scores = rng.normal(size=(len(seg), len(seg)))
block = seg[:, None] == seg[None, :]                           # attend only within a sample
attn = softmax(np.where(block, scores, -np.inf), axis=-1)
assert np.allclose(attn[~block], 0.0)                          # zero mass across sample boundaries
assert np.allclose(attn.sum(axis=-1), 1.0)                     # still a valid distribution
assert softmax(scores, axis=-1)[~block].max() > 1e-3           # adversarial: no mask -> attention leaks

# ---------- pi0.5 visual-token pruning: drop a redundant camera view ----------
# Half the tokens are visual; dropping half of those cuts attention (quadratic) by 1-(1-d)^2.
kept = 1 - 0.5 * 0.5
assert abs((1 - kept * kept) - 0.4375) < 1e-9                  # ~44% fewer attention FLOPs

print("PASS:",
      f"save@3%={s_lo*100:.2f}% save@90%={s_hi*100:.2f}% quad@90%={saving_quad*100:.2f}% |"
      f" pack {len(samples)}->{len(bins)} pad {pad_before*100:.0f}%->{pad_after*100:.1f}% |"
      f" 1.88x=>{(1-1/tput)*100:.2f}% | prune -43.8%")

How it works¶

Varlen FlashAttention keeps the fused softmax-times-value kernel of FlashAttention but changes the iteration space: instead of tiling an L x L matrix, it walks each sequence's own [start, end) range from cu_seqlens, so query block i only ever sees key blocks within the same sequence. Attention cost becomes the sum of per-sequence len_i^2 terms rather than batch * L^2. The core-math block models this with a linear (projection/MLP) plus quadratic (attention) cost and shows the two regimes the report observed: at moderate lengths the projection term dominates and savings equal the padding fraction; at long lengths the quadratic term dominates and savings exceed it.

Data packing runs first-fit-decreasing bin packing at preprocessing time: sort samples longest-first, and drop each into the first sequence with room, opening a new sequence only when none fits. A sample longer than the context is a hard error (it must be split or truncated upstream). The packed sequence carries a per-token segment id; the trainer builds a block-diagonal mask from it (seg[i] == seg[j]) so attention stays within each original sample. The core-math block proves this mask zeroes all cross-sample probability mass, and that omitting it lets attention leak, which is the single most important correctness detail.

How to use it¶

Reference template for the varlen call on a packed batch (needs a FlashAttention-2-capable stack; validate before production):

# REFERENCE TEMPLATE (needs flash-attn + torch; not executed here). Core math is validated above.
import torch
from flash_attn import flash_attn_varlen_func   # pin flash-attn; verify the signature for your version

def varlen_attention(qkv_packed, seqlens):
    # qkv_packed: [total_tokens, 3, n_heads, head_dim] -- all samples concatenated, NO padding.
    # seqlens:    per-sample lengths of the packed samples, in order.
    cu = torch.zeros(len(seqlens) + 1, dtype=torch.int32, device=qkv_packed.device)
    cu[1:] = torch.cumsum(torch.tensor(seqlens, device=qkv_packed.device), 0)  # cu_seqlens boundaries
    max_s = int(max(seqlens))
    q, k, v = qkv_packed.unbind(dim=1)
    return flash_attn_varlen_func(q, k, v, cu, cu, max_s, max_s, causal=True)  # attention within each sample only

Enable varlen on both towers of a VLA model. The vision encoder runs varlen over valid visual tokens; the language decoder runs varlen over the concatenated visual-plus-text sequence. The report notes the decoder picks the varlen or fixed path per batch based on the length distribution.
Reset position ids per packed sample. Each sample in a packed sequence must restart its positional encoding at 0; a global position counter across the whole packed sequence is a subtle correctness bug.
Cap packed length at the real context window, and route over-length samples to truncation or splitting rather than letting the packer fail mid-run.

How to develop with it (the π0.5 token-pruning case study)¶

The report's π0.5 optimization shows the same principle applied at the data level with model-specific priors, without touching the architecture:

Dynamic sequence padding. Instead of padding every batch to a global maximum (for example 200 tokens), compute max_length per batch from the actual inputs, so short batches are not padded to a long global bound. This is the batch-local version of the padding-elimination idea.
Prior-knowledge visual-token pruning. In LIBERO, the right-hand camera view was verified to contribute nothing to task success, so those image tokens are dropped before they enter the model, cutting visual-token count at the source. Because attention is quadratic, pruning a fraction of visual tokens cuts attention FLOPs super-linearly (the core-math block computes about a 44% reduction from dropping half of a visual half-sequence).

The measured result: per-step training time fell from 4.71s to 2.85s (39.56%), total training from 39h40m to 23h44m (40.2%), while the LIBERO Spatial task-success rate barely moved (98.4% to 98.2%, statistically insignificant). Note this measured ~40% training-time reduction is the honest number; the report's abstract itself attributes "165%" to training alone ("π0.5 Attention optimization has accelerated training by 165%"), a figure its own experiment section never reproduces or reconciles with the ~40% above. A separate sentence in the introduction recasts the same number as "training and inference efficiency," which only compounds the inconsistency rather than resolving it.

How to run it in production¶

Verify accuracy parity, not just speed. Packing and varlen are correctness-preserving only if masking and position ids are right. Run a downstream eval before and after enabling them; the report's evidence is parity (accuracy flat or slightly up), so a regression means a masking or normalization bug, not an inherent cost.
Profile which regime you are in. If varlen savings track the padding fraction, you are projection-bound; if they exceed it, you are attention-bound. This tells you whether packing (fewer, fuller sequences) or sequence-length reduction (pruning) is the higher-leverage lever.
Treat packing as part of the data-loading pipeline. Bin packing is preprocessing work; do it offline or in the loader, not on the training hot path, and make the packing deterministic for reproducible runs.
Keep the padding-rate metric on your dashboard. Padding fraction is a direct, cheap proxy for wasted compute; a rising padding rate is throughput leaking away.

How to maintain it¶

Re-validate the mask on every attention refactor. The block-diagonal / cu_seqlens boundary is the one thing that, if broken, silently corrupts training rather than crashing. Keep the core-math leakage assertion as a unit test; it is library-independent and survives flash-attn and torch upgrades.
Pin the FlashAttention version and re-check the varlen signature. The varlen interface arguments have changed across flash-attn releases; do not assume a snippet ports across versions.
Re-derive per-batch max_length after data changes. Dynamic padding depends on the input length distribution; a new dataset can change it, so do not hard-code a bound learned on an old mix.
Re-verify token-pruning priors per dataset and embodiment. The "right-hand camera is redundant" prior is specific to LIBERO; a different task or robot may make that view load-bearing. Pruning is a measured decision, not a default.

Results¶

Reported by Guo et al. (2026), Qwen2.5-VL backbone:

Technique	Metric	Result
Varlen attention	Time saving vs fixed-length, padding 3% to 90%	2.28% to 89.73%
Varlen attention	Long-sequence (>8k) latency reduction	25% to 90%
Data packing	Training throughput	1.88x ("188%")
Data packing	Total training-time reduction	46.87%
Data packing	Downstream accuracy (MMMU/MUIRBench/RealWorldQA/TextVQA/MMStar)	flat to slightly up
π0.5 pruning + dynamic padding	Per-step / total training time	39.56% / 40.2% faster
π0.5 pruning + dynamic padding	LIBERO Spatial success rate	98.4% to 98.2% (n.s.)

Failure modes¶

Missing or wrong block-diagonal mask. The worst failure: training proceeds, loss looks plausible, but packed samples attend across boundaries and the model learns from contamination. Always assert zero cross-sample attention.
Global position ids across a packed sequence. Each sample must restart positions at 0; a continuous counter injects a spurious long-range positional signal.
Over-length samples. A sample longer than the context window has no valid bin; the packer must refuse it and route it to truncation/splitting rather than failing the run.
Loss-normalization drift. Packing changes samples-per-sequence, so per-sequence vs per-token loss averaging can silently reweight examples.
Assuming savings without profiling. If your data already fills the context, varlen and packing buy little; the benefit is proportional to the padding you remove.

References¶

Guo et al., "Thousand-GPU Large-Scale Training and Optimization Recipe for AI-Native Cloud Embodied Intelligence Infrastructure" (JD Technology, 2026): https://arxiv.org/abs/2603.11101
Dao et al., FlashAttention-2 (the varlen kernel lives here): https://arxiv.org/abs/2307.08691
flash-attn (varlen interface, code): https://github.com/Dao-AILab/flash-attention
Zhao et al., Prepacking: fast prefilling and increased throughput (packing without cross-contamination): https://arxiv.org/abs/2404.09529
Qwen2.5-VL (the report's backbone): https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct