Step-Level Sequence Parallelism in Transformers

• Step-level sequence parallelism is a technique that partitions sequential steps or token groups to optimize transformer computations while reducing memory overhead and communication costs.
• It enables distributed execution by splitting the sequence dimension, achieving up to 13.7× batch size and 3.0× sequence length improvements, as demonstrated in transformer training experiments.
• Recent advances extend this approach with dynamic scheduling, topology-aware partitioning, and heterogeneity-adaptive assignments for both training and decode-time kernel optimizations.

Searching arXiv for recent and foundational papers on step-level/sequence parallelism. Step-level sequence parallelism denotes a family of parallelization techniques that expose concurrency at the granularity of sequence positions, contiguous sequence chunks, or sequential steps of an iterative computation. In contemporary transformer systems, the term most often refers to splitting the sequence-length dimension so that attention and adjacent operators execute over distributed or subdivided token ranges rather than over a fully replicated sequence; in decode-time kernels it can also refer to sequence-level splitting used purely as scheduling logic before any attention math is executed (Li et al., 2021, Font et al., 19 Mar 2026). Recent work extends the concept from exact distributed self-attention and compiler-inserted collectives to hierarchy-aware masking for packed hybrid-context data, heterogeneity-adaptive per-step assignment, and chunkwise serving-time scheduling (Gupta et al., 29 Apr 2026, Zhang et al., 29 Jun 2026, Wang et al., 2024, Li et al., 9 Nov 2025). In a broader algorithmic lineage, related work studies how “steps” of sequential programs can be re-expressed as parallel frontiers or antichains under explicit dependence constraints (Shen et al., 2022, 0810.5575).

1. Conceptual scope and relation to other parallelisms

Sequence parallelism, in its canonical transformer form, splits each input sequence of length $L$ into $N$ contiguous chunks of length $L/N$, placing chunk $i$ on device $i$. All devices share the same model parameters but only hold their local subsequence. This differs structurally from data parallelism, which splits batch $B$; tensor parallelism, which splits large weight matrices or head/hidden dimensions; and pipeline parallelism, which splits layer stacks. The distinctive systems property is that sequence parallelism reduces activation and attention costs that scale with $L$ by a factor of $N$, while requiring explicit communication to reconstruct full-context attention semantics (Li et al., 2021).

The newer literature broadens that basic definition in three directions. First, compiler systems treat sequence parallelism as a graph rewrite over high-level IR, automatically resizing activations and inserting collectives around attention. Second, topology-aware and hierarchy-aware systems combine multiple sequence-parallel regimes—such as all-to-all and ring communication, or inter-group and intra-group partitioning—to preserve exact causal attention on long and packed inputs. Third, inference systems introduce step-level sequence parallelism at decode time, where the mechanism may be a scheduler-level override that increases the number of launched CTAs by splitting the K/V sequence, even when the numerical attention kernel is otherwise unchanged (Gupta et al., 29 Apr 2026, Fang et al., 2024, Zhang et al., 29 Jun 2026, Font et al., 19 Mar 2026).

A useful consequence is that the phrase is not confined to a single communication primitive or to one deployment regime. In the cited work it can describe ring self-attention over distributed training shards, all-to-all head/sequence transpositions, JIT-generated NCCL P2P schedules for packed causal masks, dynamic per-step reassignment of sequences to heterogeneous SP groups, or CTA-level sequence splitting in a Hopper decode kernel.

2. Exact sequence-dimension partitioning in transformer training

A standard formalization starts from the transformer-layer input

$X\in\mathbb R^{B\times L\times H},$

with batch $B$, sequence length $N$0, and hidden size $N$1. Sequence parallelism decomposes

$N$2

so that each device stores only one sequence shard. In “Sequence Parallelism: Long Sequence Training from System Perspective,” each device computes local projections $N$3, $N$4, $N$5, and exact self-attention is recovered by Ring Self-Attention (RSA) in two ring phases: one for $N$6 score construction and one for $N$7 output accumulation. The ring rotates $N$8 and $N$9 across neighboring devices, concatenates partial scores into $L/N$0, normalizes with softmax, and produces the full output for shard $L/N$1. The resulting per-device memory is

$L/N$2

while the quadratic attention term is reduced from $L/N$3 to $L/N$4. On up to 64 NVIDIA P100 GPUs, the paper reports $L/N$5 maximum batch size and $L/N$6 maximum sequence length relative to tensor parallelism, and with sparse attention it handles over 114K tokens, over $L/N$7 longer than existing sparse-attention work holding the whole sequence on a single device (Li et al., 2021).

AutoSP moves this step-level sequence partitioning into the PyTorch-2.0 compiler. Its SP pass operates on Torch-IR, where layer-level operators are still visible and the sequence axis is not yet obscured by lower-level reshapes and permutes. The pass reads the first input tensor to infer $L/N$8 and $L/N$9, scans for an Attention node producing shape $i$0, resizes pointwise-across-sequence activations from $i$1 to $i$2, and rewrites attention activations from $i$3 to $i$4 after an inserted all-to-all. Before attention it inserts

$i$5

so attention sees full sequence length $i$6 but only $i$7 heads per device; after attention it inserts the inverse all-to-all to restore token sharding. AutoSP further modifies PyTorch-2.0’s automated activation-checkpointing by removing infinite-capacity edges on non-attention matmuls, allowing rematerialization when $i$8. In the reported experiments on GH200-96 GB and AMD MI250-64 GB, maximum trainable context increased up to $i$9 on NVIDIA and $i$0 on AMD over competitive hand-written baselines, while end-to-end throughput at 40K tokens remained near the hand-written SP baseline: $i$1 s versus $i$2 s on GH200 and $i$3 s versus $i$4 s on MI250 (Gupta et al., 29 Apr 2026).

These systems establish the core invariant of exact sequence-dimension parallelization: each device holds only a shard of tokens for most operators, but attention must still recover full-sequence semantics, either by rotating K/V blocks in a ring or by temporarily exchanging token/head ownership through all-to-all collectives.

3. Topology-aware, hierarchical, and hybrid-context extensions

A major systems question is how to preserve the exactness of attention while matching hardware topology and model structure. USP addresses this by treating all-to-all-based Ulysses and P2P-based Ring attention as orthogonal axes of a 2D process mesh of shape $i$5. Its Unified_SP_Attention first redistributes $i$6, $i$7, and $i$8 across Ulysses groups with AllToAll4D, then runs LoadBalanceRingAttention within ring groups, and finally performs the reverse all-to-all to restore the original partition. In this formulation, pure Ulysses and pure Ring are limiting cases. The abstract reports $i$9 MFU on two 8xA800 nodes for LLAMA3-8B training at sequence length 208K (Fang et al., 2024).

DSP generalizes the sharding decision to multi-dimensional transformers whose attention alternates among multiple sequence dimensions $B$0. At stage $B$1, the system shards along the dimension $B$2 used by that stage, and when the desired split dimension changes it performs a single AlltoAll to reshard. The paper’s communication model summarizes DeepSpeed-Ulysses as $B$3 total activation volume per block and DSP as $B$4, i.e. a $B$5 reduction relative to the best single-dimension method. The abstract reports throughput improvements ranging from $B$6 to $B$7, with less than $B$8 communication volume (Zhao et al., 2024).

HSAP is aimed at hybrid-context packed sequences, where naive sequence parallelism can violate causal masking by permitting cross-segment attention. Its construction splits $B$9 along sequence length into $L$0 inter-group chunks and then partitions heads within each chunk across $L$1 via all-to-all. Exact causality is preserved because attention is computed only within each original packed segment, using JIT-compiled masks that impose lower-triangular structure and zero out any cross-segment positions, for example through cu_len cumulative-length arrays. SAP’s JIT hook derives attention-block maps, balances regions, resolves conflicts globally, and emits per-rank instruction lists containing operations such as create $L$2recv-op {kv} %{src}, copy-kv comm→comp, and computing. At the HSAP level, inter-group P2P exchanges partial $L$3, intra-group all-to-all redistributes heads, and the full output is reconstructed by summing partial $L$4 without duplication. The reported results include scalability to 64 GPUs at 32K sequence-per-GPU, linear continuation to 512K sequence length on 64 GPUs with only $L$5 throughput drop, and per-GPU peak activation memory that remains near constant at approximately $L$6 GB on A100 40GB as $L$7K (Zhang et al., 29 Jun 2026).

Together, these variants show that sequence parallelism is not a monolithic scheme. The implementation space spans 2D meshes, dynamic switching among sequence axes, and hierarchy-aware decomposition designed specifically to preserve causal masks on packed data while keeping communication bounded.

4. Adaptive per-step assignment for heterogeneous training and online serving

Static sequence-parallel degrees assume homogeneous sequence lengths. FlexSP targets the long-tail distributions observed in LLM training corpora by solving a per-step assignment problem over heterogeneous SP groups. For a training step with packed sequences $L$8 of lengths $L$9, $N$0 GPUs, and up to $N$1 candidate SP groups, it introduces a binary group-selection vector $N$2, assignment matrix $N$3, and a makespan variable $N$4, then solves the MILP

$N$5

subject to runtime, memory, group-budget, and assignment constraints. The profiled runtime is

$N$6

with $N$7, and memory

$N$8

The system reduces problem size by sequence-length bucketing and, when necessary, memory-balanced micro-batch chunking via dynamic programming, then uses SCIP on the reduced MILP. On the execution side it hot-switches NCCL communicators for the chosen groups and uses flash-attn-2 for variable length. Reported speedups are up to $N$9 over DeepSpeed and $X\in\mathbb R^{B\times L\times H},$0 over Megatron-LM, with case studies showing SP=64 for very long sequences, SP=32 or SP=16 for mid-lengths, and SP=8, SP=4, or SP=1 for abundant short sequences; the profiled All-to-All fraction drops from approximately $X\in\mathbb R^{B\times L\times H},$1 to approximately $X\in\mathbb R^{B\times L\times H},$2 (Wang et al., 2024).

Serving systems face a different heterogeneity problem: a single request may require different SP degrees across its own prompt. CDSP addresses this by partitioning a request of total prompt length $X\in\mathbb R^{B\times L\times H},$3 into $X\in\mathbb R^{B\times L\times H},$4 contiguous chunks of lengths $X\in\mathbb R^{B\times L\times H},$5, assigning chunk $X\in\mathbb R^{B\times L\times H},$6 an SP group $X\in\mathbb R^{B\times L\times H},$7 of size $X\in\mathbb R^{B\times L\times H},$8, and minimizing the end-to-end prefill time through a fitted latency model

$X\in\mathbb R^{B\times L\times H},$9

Its recursive scheduler compares a single-chunk baseline against multi-chunk plans, uses SolveChunkSize to determine how many tokens should run before expanding the group size, and extends groups by selecting the least-busy instances on the same node first. The system is integrated into a disaggregated prefill/decoding cluster: prefill GPUs form an SP pool, chunk-level KV-cache is transferred to a decoding GPU through a handshake protocol, and the request is then inserted into a continuous batch chosen by Llumnix-style virtual usage. The abstract reports up to $B$0 lower TTFT under max sustainable loads, median TBT reduced by up to $B$1, and max request capacity increased by up to $B$2; the detailed ablations report per-request scheduling overhead of $B$3–$B$4 on average and $B$5 at maximum (Li et al., 9 Nov 2025).

A recurrent systems implication is that “the SP degree” is not necessarily a fixed job-level hyperparameter. In both training and serving, recent systems elevate it to a per-step or per-chunk optimization variable governed by communication cost, memory headroom, and instantaneous resource fragmentation.

5. Decode-time kernel scheduling and the sequence-aware split heuristic

In low-head-count autoregressive decoding, step-level sequence parallelism can occur entirely inside kernel-dispatch logic. FlashAttention-3 originally computed

$B$6

and applied a hard guard: if $B$7 (equivalently $B$8), force $B$9; otherwise run the usual occupancy-sweep loop. On an H100 with $N$00 SMs, that rule produces a blind spot in low-head MQA/GQA decoding, because a case such as batch $N$01, $N$02, $N$03 launches only $N$04 CTAs, leaving $N$05 SMs idle. The sequence-aware split policy changes only the boundary case $N$06, where $N$07 and $N$08. Its decision rule is: $N$09 The occupancy proxy is

$N$10

This logic is evaluated before any attention arithmetic; once $N$11 is fixed, the kernel launches $N$12 CTAs, each handling one $N$13 tile and processing $N$14 tokens (Font et al., 19 Mar 2026).

Raising $N$15 increases CTA count by $N$16, improves SM residency, and helps hide memory latency. Each CTA works on a shorter subsequence $N$17, loading its K/V slice into shared memory or reading it directly from global memory. Because the selected override is small ($N$18), the paper states that reuse remains good and the shared-memory budget is not exceeded. The cost is an additional reduction across CTAs, implemented by atomic adds on global memory or by a small tree reduction via shared memory or warp intrinsics. In the override regime, the paper characterizes that extra cost as small and more than offset by the $N$19 increase in CTAs.

For the metadata-enabled path, kernel time was measured under CUDA Graph replay and A/B interleaved timing for batch $N$20, $N$21:

$N$22	$N$23	gain / pct_gain
1	$N$24	$N$25 / $N$26
2	$N$27	$N$28 / $N$29
8	$N$30	$N$31 / $N$32

A split sweep over $N$33 at $N$34 shows that latency drops once $N$35 into a plateau near $N$36–$N$37, with the best point at $N$38 around $N$39; $N$40 was selected because it is the smallest split that captures approximately $N$41 of the benefit. Shorter contexts $N$42 are unchanged, heavier workloads $N$43 or longer $N$44 remain on the original split sweep, and a 160-config sweep reported no regression beyond $N$45 (Font et al., 19 Mar 2026).

This decode-time result is notable because it inverts a common intuition: the benefit arises not from a longer context per se, but from under-occupancy at the kernel level when head count is too small to fill Hopper SMs.

6. Algorithmic lineage, dependence structures, and recurrent misconceptions

Outside transformer systems, the underlying problem is older: how to expose parallelism among steps that appear sequential. In “Detection of parallel steps in programs with arrays,” a loop is transformed into a separated form whose body is a sequence of controllers $N$46 and a final kernel $N$47. The controllers compute indices used later; the kernel performs the data update. The paper states that any loop may be algorithmically represented in this form and that the number of controllers is invariant. Immediate data dependence between two instruction instances reduces to solving an integer connection equation

$N$48

The same work further states that if index expressions are more complex than cubical, dependence detection is algorithmically unsolvable, and specifically that polynomial indices of total degree $N$49 yield undecidability via Diophantine reduction. Its proposed escape hatch is predecessor semantics, where the right-hand side explicitly names the iteration from which data is read, trivializing dependence detection and enabling parallel execution of antichains such as hyperplanes $N$50 in stencil-style updates (0810.5575).

The phase-parallel framework of “Many Sequential Iterative Algorithms Can Be Parallel and (Nearly) Work-efficient” recasts the same issue in DAG terms. Each object $N$51 is assigned a rank $N$52, interpreted as the earliest phase in which it can be processed. All objects of one rank form a frontier

$N$53

which can be processed in parallel. The paper distinguishes work-efficiency from round-efficiency, proposes Type 1 algorithms that extract each frontier by range queries and Type 2 algorithms that wake an object when its last predecessor completes, and gives detailed instantiations for activity selection, unlimited knapsack, LIS, and greedy MIS. For LIS, the Type 2 random-pivot method achieves $N$54 work and $N$55 span; implementations in C++/Cilk Plus on a 96-core machine show 10’s–100× speedups when dependence depth is reasonable (Shen et al., 2022).

This broader lineage suggests that step-level sequence parallelism is best understood as a dependence-management pattern rather than as a single transformer kernel trick. A recurring misconception is that it always means one static scattering strategy applied uniformly to all sequences. Recent work instead includes per-stage dimension switching, per-step MILP assignment across heterogeneous SP groups, and intra-request chunk-level SP expansion in serving (Zhao et al., 2024, Wang et al., 2024, Li et al., 9 Nov 2025). Another misconception is that exact causality is incompatible with aggressive sequence partitioning on packed data; HSAP explicitly constructs JIT-compiled masks and communication schedules so that only original segments attend to themselves and all partial outputs are summed without duplication (Zhang et al., 29 Jun 2026).

A plausible implication is that future systems will continue to collapse distinctions between compile-time graph rewriting, runtime communicator selection, and kernel-dispatch occupancy control. In the current literature, the same conceptual objective—extract more useful parallel work from a sequence-ordered computation—already appears at every layer of the stack, from dependence equations and frontier schedules to NCCL plans, process meshes, and Hopper CTA grids.