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

An Efficient Hybrid Sparse Attention with CPU-GPU Parallelism for Long-Context Inference
May 8, 2026Feiyu Yao, Zhixiong Niu, Xiaqing Li +3GPU Acceleration Self-Attention
Beijing University of Technology Beijing, China · Microsoft Research Beijing, China · Beijing Jiaotong University Beijing, China
Reaching Every Position Without Searching: Rotating Sparse Wiring on the Hypercube as a Substitute for Attention
Sep 16, 2026Yoshiaki TakashitaEfficient Attention Language Modeling
School of Law, Waseda University
Morphing into Hybrid Attention Models
Jun 29, 2026Disen Lan, Jianbin Zheng, Yuxi Ren +5Efficient Attention Neural Architecture Search
1Fudan University · 2ByteDance Seed · Work done at ByteDance Seed +1

May 8, 2026J/K move · Enter open · S save

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

Feiyu Yao, Zhixiong Niu, Xiaqing Li, Yongqiang Xiong+2

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

Open paper AskPDF arXiv

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.

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