Layerwise Repacking: Variable-Length Hybrid-Attention Training without an Extra Redistribution Pass

Jiahao Zhang; Junwei Feng; Dongliang Huang; Xiaoling Zhou; Yiqiao Liao; Chengru Song

cs.LGOct 26, 2026

Layerwise Repacking: Variable-Length Hybrid-Attention Training without an Extra Redistribution Pass

Authors: Jiahao Zhang, Junwei Feng, Dongliang Huang, Xiaoling Zhou, Yiqiao Liao, Chengru Song

Organizations: KlingAI Research Beijing, China

Abstract

Variable-length training creates conflicting parallelism requirements across the layers of a hybrid-attention model. The longest samples require wide sequence-parallel (SP) groups of 16 or 32 ranks for full attention, whose cost grows quadratically with sequence length. Sparse-attention layers, dominated by near-linear MoE work, can instead balance most evaluated workloads within node-local SP8 groups. Reusing the wide groups for these layers pushes their all-to-all communication across nodes, where the observed effective rate is 7.0x lower. Yet switching to smaller groups changes which ranks own each sample; a separate redistribution at every full-sparse boundary would introduce another data transfer. Layerwise Repacking keeps each SP16/32 parent for full attention and repacks its samples into node-local sparse children at SP{1, 2, 4, 8}, where hidden states persist between full layers. The layout change rides on the two existing full-layer Ulysses sites: at each, one variable-size all-to-all performs both the head exchange and the sample repacking, so no redistribution pass is added. Because the balance the children can reach depends on which samples share a parent, an auxiliary parent solver rebalances parent membership before the children are formed. On all 22 steps with SP16/32 parents among the first 200 training steps of a 1F3S hybrid-attention model (one full-attention layer per three sparse-attention layers), repacking alone reduces the mean end-to-end time by 5.95% and raises model FLOPs utilization (MFU) by 3.75 percentage points over the baseline with repacking disabled; with the improved solver, end-to-end time falls by 7.61% and MFU rises by 4.83 percentage points.

Figures & tables

Figure 1. Overview of Layerwise Repacking . (a) Hidden states use the node-local sparse layout

\mathcal{S}

between full-attention layers; transitions to the wide layout

\mathcal{F}

are fused with the existing Ulysses exchanges. (b) An SP32 full-attention parent spans four nodes, while its sparse-layer children remain node-local. (c) Planning assigns samples to parents, optionally rebalances parent membership, and then forms sparse-layer children; dashed arrows indicate fallback paths. Three-panel diagram. The layer view shows hidden states staying in the node-local sparse layout and visiting the wide full layout only at full-attention layers. The physical view shows one SP32 parent over four nodes split into node-local children. The planning view shows rounds 1A, 1B, and 2 with their fallback paths.

Figure 2. Ulysses communication on wide and node-local ranks for the evaluated workloads. Unless noted otherwise, measurements use the baseline solver with repacking disabled. (a) Exposed all-to-all time per rank and step is 7.97 s for SP16/32 ranks and 1.60 s for node-local SP8 ranks. (b) The corresponding effective rates are 39.5 and 275.6 GiB/s. Both measures include full-layer exchanges. (c) Wide-rank payload by layer type; the full-layer share is estimated from the conversion payload measured with repacking enabled.

Stage	Fixed	Optimized output
Round-1A baseline	DP world, allowed SP	Seed topology $(G_{b},R_{b})$ and membership $I_{b}^{0}$
Round-1B improved	All $(G_{b},R_{b})$ from round-1A	Global parent membership $I_{b}^{\star}$
Repair	All $(G_{b},R_{b})$	Feasible membership $\widetilde{I}_{b}$ ( $\|\widetilde{I}_{b}\|\geq m_{b}$ )
Round-2 child planner	$(G_{b},R_{b})$ and $\widetilde{I}_{b}$	Per-parent child groups $(G_{bj},r_{bj},I_{bj})$

Table 1. Planner stages and optimization domains.

Cell	Configuration	Mean E2E (s)	Mean MFU (%)	Fwd max (s)	Bwd max (s)
B0	baseline solver / OFF (original)	37.221	56.47	8.864	23.316
B1	baseline solver / ON	35.007	60.22	7.437	22.302
N0	improved solver / OFF	35.887	58.52	8.241	22.766
N1	improved solver / ON (complete)	34.390	61.30	7.218	21.900

Table 2. Performance on all 22 SP16/32 workloads. Values average the 22 workloads; B0

\rightarrow

N1 is the complete-system comparison. All reported changes are calculated from unrounded means.

Figure 3. Results on the 22 SP16/32 workloads. (a, b) Paired E2E reduction and MFU gain for the two repacking contrasts (B0

\rightarrow

B1, N0

\rightarrow

N1), the two solver contrasts (B0

\rightarrow

N0, B1

\rightarrow

N1), and the complete system (B0

\rightarrow

N1); bars are means over the 22 workloads and whiskers are workload-level 95% bootstrap confidence intervals. (c) E2E reduction from repacking on each workload, relative to the same solver without repacking. Left two panels: bar charts with confidence intervals; E2E falls by 2.21 and 1.50 s with repacking, 1.33 and 0.62 s with the improved solver, and 2.83 s with both; MFU rises by 3.75, 2.78, 2.05, 1.08, and 4.83 percentage points. Right panel: per-workload E2E reduction from repacking, positive on 20 workloads and negative on the two-SP16 and SP32 workloads under both solvers.

Figure 4. Communication on SP16/32 ranks before and after repacking. “Original” is the parent-wide Ulysses path without repacking; “child” and “conversion” are the node-local sparse-layer exchanges and the fused full-layer exchanges. Exposed A2A time falls from 7.97/7.59 s to 4.76/4.60 s under the baseline/improved solvers. Total payload changes by less than 5%, and more than 70% of the repacked payload uses the node-local child path. Stacked bars of all-to-all time and payload on wide ranks before and after repacking under both solvers.

Metric	Baseline	Improved
Conversion payload (GiB)	74.6	70.5
Pass moving attention tensors (s)	$\approx$ 1.89	$\approx$ 1.78
Pass moving hidden state (s)	$\approx$ 0.94	$\approx$ 0.89
A2A reduction from repacking (s)	3.20	2.99
Fraction of reduction lost to the pass	30–59%	30–60%

Table 3. Estimated redistribution overhead avoided by communication fusion, per wide rank and step (forward, recomputation, and backward combined). Conversion payload and A2A time savings are measured. Estimates assume four attention-tensor transfers (Q/K/V in; output out) or two hidden-state transfers (in/out) per full layer and phase at the measured cross-node effective rate of 39.5 GiB/s, excluding additional launch and synchronization overheads. The fraction is estimated overhead divided by measured A2A time savings.

Figure 5. Where the improved solver’s gain lands: rank-max reduction in SP16/32 parent and ordinary-SP8 regions, equally averaging the repacking-OFF and -ON deltas over the same 23 wide parents. Grouped bars: forward rank max falls by 0.41 s in SP16/32 parents and rises by 0.07 s in ordinary SP8 regions; backward rank max falls by 0.37 and 0.39 s.

Communication-Efficient Distributed Training

Communication-efficient distributed training reduces the data exchanged among devices so models can be trained faster and over limited-bandwidth networks.

See all 228 papers

May 8, 2026

An Efficient Hybrid Sparse Attention with CPU-GPU Parallelism for Long-Context Inference

Long-context inference increasingly operates over CPU-resident KV caches, either because decoding-time KV states exceed GPU memory capacity or because disaggregated prefill-decode systems place KV data in host memory. Although block-sparse attention reduces attention cost in this setting, sparsity alone is insufficient for end-to-end efficiency. GPU-only designs remain constrained by PCIe bandwidth and metadata memory overhead, while CPU-GPU hybrid designs still suffer from substantial GPU idle time and bottlenecks in CPU-side top-k selection and sparse attention computation. Fluxion is built on three key insights: output-aware KV budgeting, head-specific and granularity-aware sparse configuration, and cross-device coordinated execution for sparse attention over CPU-resident KV caches. Guided by these insights, Fluxion combines a lightweight head-property predictor, a granularity-budget selector, and a priority-based scheduler to jointly optimize budget allocation, sparse configuration, and CPU-GPU execution overlap. This co-design enables hybrid sparse attention to achieve both accuracy and system efficiency in long-context inference. Across 2 models, 3 benchmarks, and 40 tasks, Fluxion preserves quality well -- the worst average degradation is only -0.26 relative to FULL, while delivering 1.5

\times

-3.7

\times

speedup over the strongest fixed sparse hybrid baseline, whose KV budget is only 0.05.

Feiyu Yao, Zhixiong Niu, Xiaqing Li +3

Beijing University of Technology Beijing, China · Microsoft Research Beijing, China · Beijing Jiaotong University Beijing, China

GPU Acceleration Self-Attention Long-Context Language Model Inference KV-Cache Management LLM Inference Acceleration Heterogeneous Computing+2

Ask PDF

Sep 16, 2026

Reaching Every Position Without Searching: Rotating Sparse Wiring on the Hypercube as a Substitute for Attention

Attention pays, at every layer and for every input, the cost of searching for whom to connect. We ask how far one can get with wiring that is fixed, sparse, and simply rotated from layer to layer. Treating the

n

positions of a sequence as the vertices of a

\log_2 n

-dimensional hypercube and connecting each position, at layer

\ell

, to its neighbour along dimension

\ell \bmod \log_2 n

, information from every position reaches every other in

\log_2 n

layers with

2n

links per layer instead of

n^2

. On a synthetic task that is unsolvable unless all positions are reached, this rotation matches all-to-all wiring at

1/32

of the links, while the same sparse pattern held fixed across layers fails; what matters is that every dimension is touched, not the order. On character-level language modelling of a public corpus (the first

12

M characters of enwik8), a hybrid that keeps two attention layers among sixteen sparse ones reaches

0.06

bits-per-character lower held-out loss than a fully attentive model of the same width at the same step budget (three seeds each, no overlap), with

1/7

of the links,

42\%

fewer parameters, and

2.4\times

less wall-clock time; the purely rotated schedule is level with the hybrid. The same ordering holds on a second corpus of mixed Japanese, English and code, where the gap widens to

0.16

. The usable learning-rate window is four to eight times wider than attention's on both. We also report what did not work - learned coordinates, and a "dynamics" variant whose apparent gains turned out to be an artefact of a saturated kernel - and the measurement discipline (frozen corpus, full-coverage evaluation, seed spread as the bar for ranking) that we found necessary to say anything at all at this scale.

Yoshiaki Takashita

School of Law, Waseda University

Efficient Attention Language Modeling Sparse Attention

Ask PDF

Jun 29, 2026

Morphing into Hybrid Attention Models

Hybrid attention models improve long-context efficiency by retaining only a subset of full-attention layers and replacing the remaining layers with linear attention. However, the effectiveness of Transformer-to-hybrid conversion critically depends on which layers preserve full attention. Existing hybrid layer selection methods typically rely on heuristic strategies such as fixed placement patterns or layerwise scoring, implicitly treating layer importance as isolated and overlooking the interdependent layer effect under a global hybrid configuration. In this work, we formulate hybrid layer selection as a budget-constrained subset optimization problem. We further propose FlashMorph (Fast LAyer Selection for Hybrid MORPHing), an effective, efficient and scalable layer selection method for Transformer-to-hybrid conversion. FlashMorph first constructs a morphable model by equipping each full-attention layer with a converted linear-attention branch. It then freezes all model weights and jointly optimizes layerwise gates on synthetic long-context retrieval data, with a linearization regularization that encourages the model to rely on linear attention for efficiency. The learned gates are discretized under a preset full-attention budget to instantiate the hybrid architecture, followed by standard logits distillation and long-context finetuning. Extensive experiments show that FlashMorph discovers more effective hybrid configurations, preserves strong long-context recall and general benchmark performance while substantially reducing layer selection cost compared with existing layer selection methods, demonstrating its effectiveness, efficiency, and scalability.

Disen Lan, Jianbin Zheng, Yuxi Ren +5

1Fudan University · 2ByteDance Seed · Work done at ByteDance Seed +1

Efficient Attention Neural Architecture Search Hybrid Transformer Architectures Hybrid Attention

Ask PDF

Abstract

Communication-Efficient Distributed Training

Explore similar work

An Efficient Hybrid Sparse Attention with CPU-GPU Parallelism for Long-Context Inference

Reaching Every Position Without Searching: Rotating Sparse Wiring on the Hypercube as a Substitute for Attention

Morphing into Hybrid Attention Models