arXiv:2604.24715v1 [cs.CL] 27 Apr 2026
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling Parsa Ashrafi Fashi ♣ 1 Utkarsh Saxena ♣ 1 Mehdi Rezagholizadeh ♣ 1 Aref Jafari 1 Akash Haridas 1 Mingyu Yang 1 Vansh Bhatia 1 Guihong Li 1 Vikram Appia 1 Emad Barsoum 1 1 AMD Correspondence to: {parsa.fashi, utkarsh.saxena, mehdi.rezagholizadeh, aref.jafari}@amd.com.
Abstract Hybrid sequence models that combine efficient Transformer components with linear sequence modeling blocks are a promising alternative to pure Transformers, but most are still pretrained from scratch and therefore fail to reuse existing Transformer checkpoints. We study upcycling as a practical path to convert pretrained Transformer LLMs into hybrid architectures while preserving short-context quality and improving long-context capability. We call our solution HyLo (HYbrid LOng-context): a long-context upcycling recipe that combines architectural adaptation with efficient Transformer blocks, Multi-Head Latent Attention (MLA), and linear blocks (Mamba2 or Gated DeltaNet), together with staged long-context training and teacher-guided distillation for stable optimization. HyLo extends usable context length by up to 32× through efficient post-training and reduces KV-cache memory by more than 90%, enabling up to 2M-token prefill and decoding in our vLLM inference stack, while comparable Llama baselines run out of memory beyond 64K context. Across 1B- and 3B-scale settings (Llama- and Qwen-based variants), HyLo delivers consistently strong short- and long-context performance and significantly outperforms state-of-the-art upcycled hybrid baselines on long-context evaluations such as RULER. Notably, at similar scale, HyLo-Qwen-1.7B trained on only 10B tokens significantly outperforms JetNemotron (trained on 400B tokens) on GSM8K, Lm-Harness common sense reasoning and RULER-64K.
1
Introduction
Transformer-based large language models (LLMs) have achieved remarkable success across a broad spectrum of tasks, including natural language understanding, reasoning, and code generation [44, 5, 8]. These advances have been driven by scaling both model size and training data, resulting in state-of-the-art performance but at the cost of substantial computational and financial resources. Consequently, training new models from scratch has become increasingly prohibitive, motivating the search for more efficient architectures and training paradigms. Recently, hybrid architectures that combine attention mechanisms with more efficient sequence modeling components such as state space models or linear attention have emerged as a promising direction. These models aim to retain the expressive power of Transformers while improving computational efficiency, particularly for long sequences. Notable examples include Jamba [28], Samba [39], Qwen3-Next [37], and Kimi-Linear [43], which demonstrate competitive performance with improved efficiency. However, these approaches largely rely on training from scratch, effectively replicating the immense cost associated with developing Transformer-based LLMs. To address this limitation, a growing line of work explores model upcycling, which seeks to convert existing pre-trained Transformer models into hybrid architectures without discarding their learned knowledge. Instead of training a hybrid model from scratch, upcycling methods reuse the parameters of a pre-trained Transformer and transform its architecture, followed by continued training. The central goal is to enable efficient knowledge transfer from a source model to a target hybrid model, thereby reducing training cost while maintaining performance. ♣ Equal Contribution First Authors, with order determined randomly.
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Figure 2: Evaluation on synthetic needle in haystack benchmark demonstrates that our upcycled hybrid 4MLA12M2 model (at only 3.9% KV cache footprint) achieves comparable performance to Llama-3.2-1B and surpasses Zebra-Llama. Furthermore, finetuning at 64K sequence length surpasses performance compared to 8K sequence length showcasing the need for long context finetuning.
Several recent works have proposed different upcycling approaches, including MambaInLlama [46], Mohawk [1], Lamba [2], and Zebra Llama [50]. These methods provide initial evidence that it is possible to re-purpose Transformer models into hybrid architectures while preserving accuracy to a certain extent. However, existing upcycling approaches predominantly focus on maintaining short-context performance metrics such as perplexity or benchmark accuracy. In doing so, they largely overlook the long-context ability of modern LLMs which has become increasingly important for real-world applications, including document understanding, code completion, and multi-hop reasoning. While hybrid architectures are often motivated by their theoret- Figure 1: Short-context math performance and average ical advantages in handling long sequences, it remains RULER accuracy across 8K, 16K, 32K and 64K context unclear whether upcycled models inherit this capability lengths. HyLo models achieve competitive short context from their Transformer counterparts.In this work, we po- performance while outperforming baselines on long-context benchmark in a limited upcycling data budget. sition long-context preservation as a core objective of upcycling, alongside short-context quality. We introduce our upcycling recipe to convert pretrained Transformer checkpoints into our HYbrid LOng-context models named HyLo without costly pretraining from scratch. Our main contributions are: • Long-context-aware model upcycling. We propose an improved upcycling recipe based on Zebra-Llama [50] yielding superior long-context performance while having comparable short-context performance (Figure 1,2). • Extended long-context training regime. Prior upcycling studies typically train to around 24K context. We scale staged training from 8K up to 64K tokens and systematically analyze how training sequence length affects long-context generalization. • Teacher-guided long-context distillation. We introduce teacher-guided long-context training with chunk-wise KL supervision, demonstrating significant gains in long-context performance while clarifying the optimization constraints introduced by this distillation design. • High throughput inference serving. We integrate HyLo into vLLM [22], enabling efficient long-context serving with tensor parallelism. HyLo enables serving contexts upto 2M tokens (30× extension over Llama-3.2-3B) on 8 AMD MI300X GPUs.
2
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
2
Related Work
Hybrid long-context models trained from scratch. Recent work has explored training hybrid architectures from scratch that combine softmax attention with more efficient sequence modeling primitives such as state-space models (SSMs) or linear attention to overcome the quadratic cost of attention. Foundational approaches include S4 [17] and Mamba [16], as well as alternative long-context mechanisms such as RetNet [42], Hyena [34], and linear attention variants [21]. Building on these ideas, recent large-scale hybrids explicitly interleave attention with efficient modules: Jamba [29] combines Transformer, Mamba, and MoE layers; Zamba [15] and Samba [40] integrate Mamba with shared or local attention; MiniMax-01 employs Lightning Attention [35] within a hybrid architecture to enable extreme context scaling [24]; Kimi Linear [43] interleaves Kimi Delta Attention with Multi-Head Latent Attention (MLA) [30]; Qwen3-Next [37] and Qwen3.5 [38] interleave softmax attention with Gated DeltaNet (GDN) [52] layers. Concurrent work also highlights key design principles: unified positional encoding across attention and SSM components is critical for stability [48], and empirical analyses show that hybrid performance depends heavily on layer allocation, gating, and memory dynamics [45], indicating that hybridization requires careful architectural co-design. Post-training upcycling and hybridization. An alternative line of work focuses on converting pretrained Transformer models into hybrid long-context models, significantly reducing training cost. Early work MambainLlama [46] shows that pretrained attention layers can initialize SSM blocks, and that retaining a subset of attention layers preserves model quality while enabling length extrapolation. This paradigm is extended by approaches such as Llamba [2] and X-EcoMLA [26], where the latter converts pretrained Transformers into MLA hybrids to improve efficiency and reduce KV-cache overhead. Closely related, Zebra-Llama [50] combines Mamba2 with MLA and introduces improved initialization, intermediate-layer distillation, and layer selection strategies, achieving near-Transformer performance with limited post-training. Additionally, L2A [7] converts softmax attention into sliding window and dynamic full attention hybrid. Subsequent work focuses on identifying which attention components are essential during conversion: methods such as RAD detect redundant attention layers [19], KL-guided approaches optimize hybrid layer allocation [27], and HALO/HypeNet improve positional adaptation under constrained budgets [6], while retrieval-aware distillation shows that preserving only a small subset of retrieval-critical attention heads is sufficient to maintain long-context reasoning performance [3]. Beyond standard language modeling, hybridization has also been explored for reasoning efficiency: the M1 model introduces a hybrid Mamba-based architecture trained with distillation and reinforcement learning, demonstrating that hybrid designs can achieve competitive reasoning performance with improved inference efficiency [47].
3
Methodology
Our goal is to upcycle pretrained Transformer LLMs into long-context hybrid models while preserving short-context quality. To this end, we propose HyLo, an efficient training recipe that extends context length and improves long-range modeling without pretraining from scratch. Building on MambaInLlama [46] and Zebra-Llama [50], which show that careful initialization and staged distillation preserve short-context performance, our method treats long-context preservation as a first-class objective. Our key contributions are a stronger architecture recipe, staged long-context training, and a broader evaluation across model families and linear block types. Architecture Design To reduce the quadratic cost of full attention, we use a hybrid architecture that combines Multi-head Latent Attention (MLA) [30] with linear recurrent blocks, including Mamba-2 (M2) [11] and Gated DeltaNet (GDN) [52]. The MLA-to-linear ratio defines the quality–efficiency trade-off: more MLA layers increase attention capacity but also increase KV-cache usage, whereas Mamba-2 and GDN add no KV-cache overhead. Unlike prior upcycling studies that focus on one base model and one linear module, we evaluate two Transformer families (Llama and Qwen) and two linear block types (Mamba-2 and GDN), showing that the recipe generalizes across architectures and scales.
3.1
Initialization
A key challenge in upcycling is how to initialize replaced hybrid blocks from a pretrained attention-based model. 3
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Following Zebra-Llama [50], we first construct a pure MLA model and a pure linear model (Mamba-2 or GDN) by replacing all attention blocks in the base Transformer. Each pure model is then initialized from the original pretrained weights. Initialization schemes for Mamba-2 and MLA are introduced in MambaInLlama [46] and X-EcoMLA [25]. Here, we describe our procedure for initializing GDN blocks from Transformer checkpoints. Our GDN Initialization. In the GDN-based HyLo hybrid architecture, each selected decoder layer replaces the standard attention module with a GDN mixer, while preserving the SwiGLU MLP and RMSNorm sublayers from the original Transformer block (see Section A.3). Starting from a pretrained model, each designated GDN layer undergoes in-place module replacement. The MLP weights and RMSNorm parameters are copied verbatim from the corresponding Transformer layer. For attention-to-GDN weight transfer, we address dimension mismatches between the projection weights of the two modules: 1. Grouped-Query Attention (GQA) expansion: When the teacher uses Hkv < Hq key-value heads (e.g., 8 vs. 32 in Llama-3.2-1B), the K and V weight matrices are first expanded by repeating each KV head Hq /Hkv times: K W̃K = RepeatKV(Wteacher , g = Hq /Hkv ).
(1)
2. Dimension truncation: Since GDN’s key dimension dk < d and value dimension dv > d, we transfer the overlapping submatrices: Q Q WGDN [ : dk , :] ← Wteacher [ : dk , :], K K WGDN [ : dk , :] ← Wteacher [ : dk , :], V V WGDN [ : min(d, dv ), :] ← Wteacher [ : min(d, dv ), :],
(2)
O O WGDN [:, : min(d, dv )] ← Wteacher [:, : min(d, dv )].
GDN-specific parameters—the gate projection WG , decay parameters (Alog , ∆bias ), beta projection Wβ , and short convolution kernels—remain at their default random initialization.
3.2
Two-Stage Light Fine-Tuning
After initialization, we apply two light fine-tuning stages: (i) our enhanced intermediate layer distillation (Enhanced-ILD) and (ii) long context supervised fine-tuning (SFT). In Stage I, pure MLA/Mamba2/GDN models undergo Enhanced-ILD training on only 20% of the data. We then assemble the final hybrid model from these Stage-I checkpoints and proceed to Stage II. Stage I: Our Enhanced-ILD Training. Zebra-Llama [50] uses ILD to refine initialization by aligning per-layer hidden states. In HyLo, we extend this objective by adding an ILD term on token-mixer outputs (i.e., Transformer attention outputs and their corresponding MLA/Mamba2/GDN outputs), which shows a significant improvement to our training. Therefore, for each layer ℓ, we minimize the sum of L2 distances between teacher and student hidden states and token-mixer outputs: L h i X (s) (t) (s) (t) LILD = hℓ − hℓ 2 + aℓ − aℓ 2 , (3) ℓ=1 (s) (t) (s) (t) where hℓ , hℓ are the student and teacher hidden states after layer ℓ, and aℓ , aℓ are the corresponding attention/token-
mixer outputs. This extra ILD term strengthens knowledge transfer from full attention to MLA/Mamba2/GDN blocks; its impact is reported in Table 6. We keep the Stage-I context length fixed at 2K. Stage II: Long-Context SFT Training. In Stage II, we load the separately distilled MLA/Mamba2/GDN checkpoints from Stage I and assemble them into one hybrid model. Because our focus is long-context extension rather than layer selection, we use either uniform layer selection or baseline-recommended layouts (when available). At this stage, we extend training context length from 2K to 8K and 64K, which is another core contribution. The assembled hybrid model is then fine-tuned end-to-end with output-level knowledge distillation using KL divergence at extended context lengths: LSFT = DKL softmax z(s) softmax z(t) , (4) where z(s) and z(t) are the student and teacher logits, respectively. 4
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Config No Teacher No Teacher + FusedLinearCE + Act checkpoint No Teacher + Act checkpoint No Teacher + FusedLinearCE + Act checkpoint 8B Teacher 8B Teacher + Fused KL 8B Teacher + Act checkpointing 8B Teacher + Fused KL Hidden 8B Teacher + Fused KL + Act checkpointing 8B Teacher + Fused KL Hidden + Act checkpointing
Memory (GiB) OOM 137.9 131 29.6 OOM OOM OOM 158.8 144.8 54.2
Table 1: Training memory for upcycling a Llama-1B model with 4 MLA and 12 Mamba-2 layers at 64K context length, with teacher (KD loss) and without teacher (CE loss).
3.3
Memory-Efficient Long-Context Distillation
Extending knowledge distillation from 2K to 64K context length introduces severe memory pressure. The dominant bottleneck is the logit tensor: for sequence length T and vocabulary size V , standard KL divergence requires materializing both student and teacher logits of shape (T, V ). At T =65,536 and V =128,256 (Llama-3), each logit tensor consumes approximately 16 GB in bfloat16, making naive distillation infeasible even on 80 GB GPUs. We address this with progressively stronger memory optimizations (summarized in Table 1). Without a teacher, the 64K setup is still OOM unless we combine activation checkpointing with FusedLinearCE, which reduces memory to 29.6 GiB. With an 8B teacher, naive KD, Fused KL-only, and checkpoint-only settings remain Out of Memory (OOM); combining activation checkpointing with Fused KL Hidden reduces memory to 54.2 GiB. Together, these optimizations enable a 32× increase in training context length (2K→64K) with an 8B teacher while maintaining single-epoch training on 8 GPUs. Further implementation details and a configuration breakdown across context lengths are provided in Appendix A.4.
3.4
vLLM Runtime Integration
To enable practical deployment of HyLo, we integrate it into vLLM [22]. This requires extending the vLLM inference stack to support architectures that interleave Mamba/GDN sequence modeling layers with MLA layers, a combination not anticipated by existing serving engines. Our integration addresses three systems-level challenges: (1) execution of heterogeneous layer types (Mamba SSM, GDN linear attention and MLA attention) within a unified serving engine, where the scheduler must manage both a fixed-size Mamba hidden state and a variable-size MLA KV cache; (2) support for MLA-specific KV compression and head expansion mechanisms, which differ from standard grouped-query attention and require custom cache allocation logic; and (3) kernel limitations arising from model-specific head dimensions (e.g., HyLo’s compressed latent dimension) that are not directly supported by existing fused attention implementations such as FlashAttention, necessitating fallback to PyTorch-based kernels with associated overhead. We implement the required runtime adaptations and evaluate their impact on long-context serving efficiency under paged attention and continuous batching.
4
Experiments and Results
4.1
Experimental Setup
Model Configurations. We implement our upcycling recipe starting from three base models: Llama-3.2-1B, Llama-3.2-3B, and Qwen3-1.7B. Full model configurations and training hyperparameters are provided in Appendix (Table 7). Evaluation Tasks. We use the lm-eval-harness [12] for short context, long context and math reasoning evaluations of our model. For short context common sense reasoning we perform evaluations on language understanding tasks, 5
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Model and Setting
Teacher
KV
Common Sense Reasoning ↑
RULER ↑
cache ARC ARE HS OB PI RA WG Avg. 8K 16K 32K 64K Baseline Models MambaInLlama-1B-50% 8B 50% 37.7 65.5 58.2 37.6 73.2 36.5 59.3 52.6 18.9 3.0 1.0 0.0 Llamba 1B 1B+70B 0% 37.1 65.4 61.3 36.8 73.8 37.6 60.6 53.2 2.9 0.0 0.0 0.0 Zebra-Llama-1B (4MLA-12M2) 8B 4% 39.1 65.4 56.9 37.0 72.3 34.5 57.9 51.8 12.3 6.8 3.7 0.1 Zebra-Llama-1B (8MLA-8M2) 8B 7.8% 38.0 66.4 58.2 38.0 72.7 36.9 61.3 53.1 0.5 0 0.1 0 Training Context Length = 8K HyLo-Llama-4MLA12M2 8B 3.9% 38.1 65.7 57.6 37.0 72.5 35.4 58.6 52.1 53.1 10.6 2.0 0.5 HyLo-Llama-4MLA12GDN 8B 3.9% 38.6 66.9 59.1 37.6 72.7 36.7 60.1 53.1 55.1 11.9 2.4 0.8 HyLo-Llama-8MLA8M2 8B 7.8% 38.8 66.7 58.3 37.0 72.3 37.4 59.7 52.9 59.0 0.3 0.1 0.0 HyLo-Llama-8MLA8GDN 8B 7.8% 39.3 67.2 59.3 37.6 72.0 38.4 60.1 53.4 60.3 0.5 0.1 0.1 Training Context Length = 64K HyLo-Llama-4MLA12M2 8B 3.9% 35.7 63.3 55.3 34.8 71.3 34.7 56.8 50.3 53.3 46.7 40.4 37.9 HyLo-Llama-4MLA12GDN 8B 3.9% 36.0 63.6 57.4 38.2 70.7 35.4 57.2 51.2 52.5 48.3 44.5 40.8 HyLo-Llama-8MLA8M2 8B 7.8% 36.3 63.4 56.1 35.0 71.2 36.6 58.6 51.0 59.0 52.5 45.5 38.8 HyLo-Llama-8MLA8GDN 8B 7.8% 36.4 64.4 57.2 37.0 72.3 37.1 58.4 51.8 61.5 53.7 48.1 41.6
GSM8K ↑
16.2 12.5 37.2 43.4 49.2 51.9 51.0 54.6 33.0 37.5 40.0 39.4
Table 2: Comparison of different techniques across backbone models Llama-3.2-1B. Model and Setting
Teacher
KV
Common Sense Reasoning ↑
RULER ↑
cache ARC ARE HS OB PI RA WG Avg. 8K 16K 32K 64K Baseline Models Mamba in Llama-3B-50% 70B 50.0% 47.1 74.0 69.0 38.4 75.9 40.1 66.5 58.7 37.0 1.0 0.0 0.0 Llamba 3B 3B+70B 0.0% 45.7 73.8 73.3 42.4 78.0 40.1 70.0 60.5 3.5 0.0 0.0 0.0 M1 21.4% 45.6 72.6 61.5 39.4 73.3 35.9 64.9 56.2 63.5 43.6 30.3 17.4 Zebra-Llama 3B (6MLA-22M2) 8B 2.0% 44.7 70.8 67.7 38.8 75.6 39.4 64.5 57.4 42.5 0.4 0.5 0.3 Zebra-Llama 3B (14MLA-14M2) 8B 4.7% 45.7 71.8 68.6 38.6 75.7 40.9 64.4 58.0 35.1 13.3 6.3 4.2 Training Context Length = 8K HyLo-Llama-6MLA22M2 8B 2.0% 45.6 72.4 67.8 38.4 76.1 39.7 66.8 58.1 65.7 39.5 25.2 11.4 HyLo-Llama-6MLA22GDN 8B 2.0% 45.4 71.9 69.3 42.4 76.6 39.7 67.6 59.0 71.2 45.0 27.1 14.1 HyLo-Llama-14MLA14M2 8B 4.7% 46.3 73.0 68.7 40.4 75.9 40.4 67.7 58.9 75.3 49.7 16.6 0.4 HyLo-Llama-14MLA14GDN 8B 4.7% 47.3 72.6 69.5 40.0 76.3 41.5 67.3 59.2 71.1 45.6 19.2 0.2 Training Context Length = 64K HyLo-Llama-6MLA22M2 8B 2.0% 43.5 69.7 66.2 38.8 75.5 38.6 64.3 56.7 65.4 56.4 49.9 42.3 HyLo-Llama-6MLA22GDN 8B 2.0% 43.7 69.5 67.9 38.6 75.9 39.8 64.9 57.2 68.2 62.1 55.7 46.3 HyLo-Llama-14MLA14M2 8B 4.7% 44.1 71.2 67.3 39.6 75.4 40.0 64.3 57.4 71.7 65.4 57.8 46.6 HyLo-Llama-14MLA14GDN 8B 4.7% 45.1 72.0 68.2 39.4 76.1 40.9 63.8 57.9 73.2 69.7 62.9 52.0
GSM8K ↑
56.8 47.8 62.5 60.7 66.2 66.3 68.1 71.0 68.2 51.6 56.0 40.9 58.9
Table 3: Comparison of different techniques across backbone models Llama-3.2-3B.
which includes ARC-Challenge (ARC) [9], ARC-Easy (ARE) [9], HellaSwag (HS) [53], OpenBookQA (OB) [31], PIQA [4], RACE (RA) [23], and WinoGrande (WG) [41]. For long context evaluations we use all 13 tasks from RULER [20] benchmark. For math reasoning, we include GSM8K [10]. Baselines. We compare with hybrid model upcycling approaches including MambainLlama[46], Llamba[2], ZebraLlama[50], M1[47] and HypeNet[6]. Among these, HypeNet proposes a hybrid upcycling approach which attempts to maintain long-context performance of the teacher. Additionally, we compare our results with Jet-Nemotron-2B[18] which is pretrained on 200B tokens from scratch.
4.2
Main Results
Tables 2,3,4 present results across Llama-3.2-1B, Llama-3.2-3B, and Qwen backbone models, comparing HyLo with baseline models on Common Sense Reasoning, GSM8K, and long-context reasoning using RULER. After long-context training, we observe a small drop in short-context performance on Common Sense Reasoning benchmarks for some models, which is expected when adapting models to longer context lengths. However, the performance degradation is relatively small across all backbone models and training settings. GSM8K performance remains competitive across most configurations. In contrast, we observe substantial improvements on long-context reasoning tasks. HyLo models significantly outperform baseline models on the RULER benchmark, particularly at longer evaluation context lengths such as 32K and 64K tokens. Our models maintain much stronger performance as context length increases, suggesting that the proposed training recipe improves the model’s ability to effectively generalize over long context. 6
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Model and Setting
Teacher
Jet Nemotron-2B Hype Net (7FA21LA)
– 1.7B
HyLo-Qwen-7MLA21M2 HyLo-Qwen-7MLA21GDN HyLo-Qwen-14MLA14M2 HyLo-Qwen-14MLA14GDN
8B 8B 8B 8B
HyLo-Qwen-7MLA21M2 HyLo-Qwen-7MLA21GDN HyLo-Qwen-14MLA14M2 HyLo-Qwen-14MLA14GDN
8B 8B 8B 8B
Common Sense Reasoning ↑
KV
RULER ↑
GSM8K ↑
cache ARC ARE HS OB PI RA WG Avg. 8K 16K 32K 64K Baseline Models 2.1% 42.5 54.5 64.4 34.0 73.5 35.4 64.9 52.7 71.3 60.1 43.9 14.1 25% 41.6 67.9 57.4 36.6 72.7 32.9 63.1 53.2 36.4 31.3 23.8 16.4 Training Context Length = 8K 3.9% 44.3 71.4 60.5 39.4 73.3 36.3 61.5 55.2 58.7 41.1 27.5 14.6 3.9% 45.3 72.7 61.5 39.4 73.2 36.3 64.6 56.1 63.5 43.6 30.3 17.4 7.8% 45.9 73.2 61.2 39.2 73.8 36.5 63.1 56.1 74.2 58.6 33.5 10.7 7.8% 45.9 73.3 62.1 38.4 74.8 37.4 63.4 56.5 71.1 45.6 19.2 0.2 Training Context Length = 64K 3.9% 42.7 70.0 60.3 38.0 73.5 35.7 63.8 54.9 56.5 49.0 38.4 27.8 3.9% 44.2 71.4 61.2 37.4 73.7 36.9 63.1 55.4 59.8 53.8 42.5 30.5 7.8% 45.4 73.3 61.1 37.8 73.7 36.6 61.8 55.7 73.9 62.6 46.2 33.1 7.8% 47.9 74.6 61.9 38.2 75.0 36.8 62.3 56.7 66.9 53.2 41.4 31.6
19.4 1.1 72.3 76.0 75.8 76.1 69.9 73.3 73.5 73.8
1B-4MLA12M2
60
52.6
49.4 52.2 50.9 50.3
52.9 51.8 51.0
58.1
1B-8MLA8M2
53.9 56.7
40 20 0
2KK+Yarn 8KK+Yarn 64K 2 8 1B-4MLA12M2
8KK+Yarn 64K 8 1B-8MLA8M2
8KK+Yarn 64K 8 3B-6MLA22M2
3B-6MLA22M2
40 RULER-64K
Common Sense Reasoning
Table 4: Comparison of different techniques across backbone models QWEN.
20 0
42.3
31.3
30 10
39.6
38.8
37.9
18.3 11.4 4.9
0.5 0.1 2KK+Yarn 8KK+Yarn 64K 2 8 1B-4MLA12M2
0.0 8KK+Yarn 64K 8 1B-8MLA8M2
8KK+Yarn 64K 8 3B-6MLA22M2
Figure 3: Impact of training sequence length and position interpolation using Yarn. Applying Yarn extension improves long context performance with a slight degradation in short context commonsense reasoning abilities. Furthermore, training at longer context preserves the long context abilities to a greater extent.
We also compare models trained with different training context lengths (8K and 64K). Models trained with longer training contexts generally achieve better performance on longer evaluation contexts while showing only modest decreases in short-context reasoning performance (also observed in [13]). Overall, the results demonstrate that our long-context training recipe effectively improves long-context reasoning while maintaining strong short-context and mathematical reasoning performance across multiple backbone models.
4.3
Ablation Studies
Comparison with position interpolation. To reduce the computational cost of model upcycling, we train models with shorter context lengths and then apply zero-shot context length extension. Specifically, we train models at different sequence lengths while keeping the total training token budget constant, and then apply YaRN position interpolation [33] to the RoPE embeddings in MLA layers to extend the context length. Mamba layers, which do not use positional embeddings, remain unchanged. In Figure 3, we evaluate performance on both short- and long-context tasks. Applying YaRN slightly reduces short-context accuracy but significantly improves long-context performance. For example, the 1B-4MLA-12M2 model trained with an 8K context achieves 50.7% average accuracy on short-context tasks but only 0.5% on RULER at 64K. After YaRN scaling, short-context accuracy decreases slightly to 49.0%, while long-context performance improves to 31.3% at 64K. Importantly, we observe that training directly with a 64K context length yields the best long-context performance while maintaining comparable short-context accuracy. Similar trends are observed for the 1B-8MLA-8M2 and 3B-6MLA-22M2 models. These results demonstrate that long-context training is effective for hybrid model upcycling. Impact of knowledge distillation. We analyze the effectiveness of knowledge distillation (KD) in our long-context training recipe. While prior work (e.g., MambaInLlama [46] and Zebra Llama [50]) showed KD improves short-context performance, we study its impact on long-context learning. As shown in Figure 4, KD has a substantially larger effect on long-context performance than 7
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
50 40 30 20 10 0
8K+Yarn 45.40
48.50
50.90
64K 45.40 47.50
1bT
8bT
50.30
RULER-64K
Common Sense Reasoning
N/A
40
8K+Yarn
30
27.70
20
64K
37.90
31.30
18.10
24.20 15.40
10 N/A 1bT 8bT
0
N/A 1bT 8bT
N/A 1bT 8bT
N/A 1bT 8bT
Figure 4: Impact of size of teacher at long context knowledge distillation. Larger teacher improves both short-context common sense reasoning tasks as well as long context ability. Model and Setting
Common Sense Reasoning ↑ ARC ARE HS OB
PI
RA WG Avg. 4K
RULER ↑ 8K 16K 32K 64K
1B-4MLA12M2 36.6 64.3 55.5 35.6 71.4 35.5 57.5 49.1 50.6 44.1 41.6 38.6 31.3 1B-4MLA12M2 w/ attn. gating 37.0 63.9 54.7 34.4 70.3 35.6 57.2 48.6 52.4 44.3 41.9 38.5 29.3 1B-4MLA12M2 w/ NoPE 38.5 66.5 57.0 36.0 71.8 35.0 56.8 49.9 59.2 51.1 4.8 1.4 0.0
Table 5: Ablation on architectural design choices incorporated in our upcycled hybrid models. (a) Removes utilizes No Positional Embeddings (NoPE) in MLA layers,while (b) adds learnable gating after attention output in MLA layers. While both NoPE and attention gating have been shown to improve long context generalization, same trends do not hold for our upcyled hybrid models.
on short-context tasks. For the 1B-4MLA-12M2 model trained at 64K context length, using an 8B teacher improves short-context reasoning accuracy by 6%, while RULER accuracy at 64K improves by 22%. Larger teacher models consistently yield greater gains. When training at 8K context length and extending context using YaRN, KD still improves RULER-64K accuracy by 14%, showing KD remains effective even when long-context ability is obtained via post-training context extension. Overall, combining KD with long-context training significantly improves performance under the same training token budget. Ablation on architectural design choices. Recently, several architectural modifications have been proposed to improve long-context performance. No Position Embedding (NoPE) [49] removes positional information from full attention layers and improves extrapolation beyond the training length. DRoPE [14] further shows that pretraining with RoPE followed by finetuning with NoPE yields strong performance. Another orthogonal approach, Gated Attention [36], applies a multiplicative learnable sigmoid gate to attention outputs and has also been shown to improve long-context extrapolation. Motivated by these works, we evaluate NoPE and Gated Attention in our long-context hybrid model upcycling framework. We train models at 8K context length and extend context using YaRN. However, neither method improves performance in our setting. Applying NoPE to MLA layers yields no long-context generalization gains. Gated Attention provides small improvements at 4K, 8K, and 16K, but the gains diminish at longer contexts, and performance at 64K is 1% lower than our baseline 1B-4MLA-12M2 model trained at 8K without these modifications. These results suggest that while NoPE and Gated Attention are beneficial when included during pretraining, they do not provide improvements in the hybrid upcycling setting. Impact of our Enhanced-ILD loss. In this ablation, we evaluate the effect of our Enhanced Intermediate-Layer Distillation (Enhanced-ILD) loss, which aligns token-mixing representations between the Transformer teacher and the corresponding hybrid student blocks (MLA/M2/GDN). Table 6 shows results for models trained with the regular ILD loss introduced in Zebra-Llama [50] versus our Enhanced-ILD loss. We observe consistent improvements across model scales and hybrid compositions. On commonsense reasoning, Enhanced-ILD yields stable gains in average score: from 51.8 to 52.8 (+1.0) for 1B (4MLA-12M2), from 53.1 to 53.4 (+0.3) for 1B (8MLA-8M2), and from 62.1 to 62.3 (+0.2) for 8B (8MLA-24M2). More importantly, Enhanced-ILD provides a much larger and more consistent boost 8
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Common Sense Reasoning ↑
Model and Setting
ARC ARE HS OB 1B-4MLA12M2 1B-4MLA12M2 + Our Enhanced-ILD 1B-8MLA8M2 1B-8MLA8M2 + Our Enhanced-ILD 8B-8MLA24M2 8B-8MLA24M2+ Our Enhanced-ILD
39.1 38.7 38.0 37.5 52.1 52.1
65.4 66.7 66.4 66.9 77.1 76.9
PI
GSM8K ↑
RA WG Avg.
56.9 37.0 72.3 34.5 57.9 51.8 57.9 37.8 72.7 36.4 59.0 52.8 58.2 38.0 72.7 36.9 61.3 53.1 58.6 38.0 73.6 37.6 61.6 53.4 74.3 41.8 78.8 40.8 69.7 62.1 74.5 42.4 79.0 42.5 69.1 62.3
37.2 43.5 43.4 48.8 66.3 72.4
Table 6: Ablation of the impact of Enhanced-ILD loss.
Figure 5: TTFT and TPOT comparison for 3B models with backbone model Llama-3.2-3B on vLLM.
on GSM8K: 37.2 to 43.5 (+6.3), 43.4 to 48.8 (+5.4), and 66.3 to 72.4 (+6.1), respectively. These results indicate that Enhanced-ILD is especially effective for strengthening mathematical reasoning while preserving, or slightly improving, broad commonsense performance.
4.4
Inference Latency Evaluation
All experiments reported use vLLM with TP=8, batch size=1, on a single node with 8 AMD Instinct MI300X GPUs. Each model is tested on a context-length sweep from 8K to 2M tokens, measuring prefill and decode latency. Prefill latency. Figure 5 reports TTFT across context lengths. At 8K–64K, all three models show comparable prefill latency. Beyond 64K, Llama 3B runs OOM, as its 28 attention layers each maintain a full KV cache whose combined footprint exceeds GPU memory. Both HyLo variants complete the full sweep up to 2M context length: while HyLo-Llama-6MLA22M2 is about 2.2x faster than HyLo-Llama-14MLA14M2 at 2M, directly reflecting the O(n2 ) cost of 14 vs. 6 MLA layers. Decode latency. Figure 5 shows per-token decode latency. At short contexts (8K–32K), Llama 3B achieves lower TPOT. However, its TPOT grows linearly with context as KV cache access scales across all 28 layers, and the model OOMs at 128K. HyLo-Llama-6MLA22M2 maintains a flat TPOT from 8K through 64K, as the Mamba layers use a fixed-size hidden state rather than an expanding cache. Beyond 64K, TPOT rises sub-linearly as the MLA layers’ KV cache grows and at 2M, HyLo-Llama-6MLA22M2 reaches about 2x faster throughput compared with HyLo-Llama-14MLA14M2.
5
Conclusion
In this paper, we present HyLo, a series of hybrid LLMs upcycled from pretrained Transformer checkpoints, with explicit emphasis on preserving long-context capability. We introduced long-context aware upcycling strategy that combines MLA-based transformer attention blocks with linear blocks instantiated with both Mamba2 and GDN, staged context-length expansion and teacher guided distillation. Across 1B- and 3B-scale settings, including both Llama- and Qwen-based backbones, our results indicate that HyLo achieves superior long-context generalization while maintaining competitive short-context quality compared to related hybrid model upcycling baselines. Additionally, Beyond quality, 9
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
HyLo is deployment-oriented: our models reduce KV-cache memory by more than 90% and, with our integrated vLLM runtime, support up to 2M-token prefill and decoding. As future work, we plan to further close the remaining gap at long context lengths, improve distillation efficiency, and extend this framework to broader downstream settings where robust long-context reasoning is essential.
References [1] Aviv Bick, Kevin Y Li, Eric P Xing, J Zico Kolter, and Albert Gu. Transformers to ssms: Distilling quadratic knowledge to subquadratic models. arXiv preprint arXiv:2408.10189, 2024. [2] Aviv Bick, Tobias Katsch, Nimit Sohoni, Arjun Desai, and Albert Gu. Llamba: Scaling distilled recurrent models for efficient language processing. arXiv preprint arXiv:2502.14458, 2025. [3] Aviv Bick, Eric P Xing, and Albert Gu. Retrieval-aware distillation for transformer-ssm hybrids. arXiv preprint arXiv:2602.11374, 2026. [4] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432–7439, 2020. [5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020. [6] Yingfa Chen, Zhen Leng Thai, Zihan Zhou, Zhu Zhang, Xingyu Shen, Shuo Wang, Chaojun Xiao, Xu Han, and Zhiyuan Liu. Hybrid linear attention done right: Efficient distillation and effective architectures for extremely long contexts. arXiv preprint arXiv:2601.22156, 2026. [7] Sakshi Choudhary, Aditya Chattopadhyay, Luca Zancato, Elvis Nunez, Matthew Trager, Wei Xia, and Stefano Soatto. Learning when to attend: Conditional memory access for long-context llms, 2026. URL https: //arxiv.org/abs/2603.17484. [8] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022. [9] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457, 2018. [10] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. CoRR, abs/2110.14168, 2021. URL https://arxiv.org/abs/2110.14168. [11] Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21060, 2024. [12] Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation. 2023. 10
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
[13] Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. How to train long-context language models (effectively). In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 7376–7399, 2025. [14] Yoav Gelberg, Koshi Eguchi, Takuya Akiba, and Edoardo Cetin. Extending the context of pretrained llms by dropping their positional embeddings, 2025. URL https://arxiv.org/abs/2512.12167. [15] Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: A compact 7b ssm hybrid model. arXiv preprint arXiv:2405.16712, 2024. [16] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First conference on language modeling, 2024. [17] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396, 2021. [18] Yuxian Gu, Qinghao Hu, Shang Yang, Haocheng Xi, Junyu Chen, Song Han, and Han Cai. Jet-nemotron: Efficient language model with post neural architecture search, 2025. URL https://arxiv.org/abs/2508.15884. [19] Yuichiro Hoshino, Hideyuki Tachibana, Muneyoshi Inahara, and Hiroto Takegawa. Rad: Redundancy-aware distillation for hybrid models via self-speculative decoding. arXiv preprint arXiv:2505.22135, 2025. [20] Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What’s the real context size of your long-context language models? arXiv preprint arXiv:2404.06654, 2024. [21] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156–5165. PMLR, 2020. [22] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention, 2023. URL https://arxiv.org/abs/2309.06180. [23] Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. arXiv preprint arXiv:1704.04683, 2017. [24] Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint arXiv:2501.08313, 2025. [25] Guihong Li, Mehdi Rezagholizadeh, Mingyu Yang, Vikram Appia, and Emad Barsoum. X-ecomla: Upcycling pre-trained attention into mla for efficient and extreme kv compression. arXiv preprint arXiv:2503.11132, 2025. [26] Guihong Li, Mehdi Rezagholizadeh, Mingyu Yang, Vikram Appia, and Emad Barsoum. X-ecomla: Upcycling pre-trained attention into mla for efficient and extreme kv compression. arXiv preprint arXiv:2503.11132, 2025. [27] Yanhong Li, Songlin Yang, Shawn Tan, Mayank Mishra, Rameswar Panda, Jiawei Zhou, and Yoon Kim. Distilling to hybrid attention models via kl-guided layer selection. arXiv preprint arXiv:2512.20569, 2025. [28] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: A hybrid transformer-mamba language model, 2024. [29] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, et al. Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887, 2024. 11
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
[30] Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434, 2024. [31] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:1809.02789, 2018. [32] Maxim Milakov and Natalia Gimelshein. arXiv:1805.02867, 2018.
Online normalizer calculation for softmax.
arXiv preprint
[33] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.00071, 2023. [34] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pp. 28043–28078. PMLR, 2023. [35] Zhen Qin, Weigao Sun, Dong Li, Xuyang Shen, Weixuan Sun, and Yiran Zhong. Lightning attention-2: A free lunch for handling unlimited sequence lengths in large language models. arXiv preprint arXiv:2401.04658, 2024. [36] Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, Fei Huang, Suozhi Huang, Dayiheng Liu, Jingren Zhou, and Junyang Lin. Gated attention for large language models: Non-linearity, sparsity, and attention-sink-free, 2025. URL https://arxiv.org/abs/2505.06708. [37] Qwen Team. Qwen3-next: Towards ultimate training & inference efficiency. https://qwen.ai/blog?id= 4074cca80393150c248e508aa62983f9cb7d27cd, September 2025. Accessed: 2026-03-19. [38] Qwen Team. Qwen3.5: Towards native multimodal agents. https://qwen.ai/blog?id=qwen3.5, February 2026. Accessed: 2026-03-19. [39] Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522, 2024. [40] Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. URL https://arxiv. org/abs/2406.07522, 2406: 07522, 2024. [41] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106, 2021. [42] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023. [43] Kimi Team, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang, Zhiyuan Li, et al. Kimi linear: An expressive, efficient attention architecture. arXiv preprint arXiv:2510.26692, 2025. [44] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [45] Dustin Wang, Rui-Jie Zhu, Steven Abreu, Yong Shan, Taylor Kergan, Yuqi Pan, Yuhong Chou, Zheng Li, Ge Zhang, Wenhao Huang, et al. A systematic analysis of hybrid linear attention. arXiv preprint arXiv:2507.06457, 2025. [46] Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. Advances in Neural Information Processing Systems, 37:62432–62457, 2024. [47] Junxiong Wang, Wen-Ding Li, Daniele Paliotta, Daniel Ritter, Alexander M Rush, and Tri Dao. M1: Towards scalable test-time compute with mamba reasoning models. arXiv preprint arXiv:2504.10449, 2025. [48] Bingheng Wu, Jingze Shi, Yifan Wu, Nan Tang, and Yuyu Luo. Transxssm: A hybrid transformer state space model with unified rotary position embedding. arXiv preprint arXiv:2506.09507, 2025. 12
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
[49] Bowen Yang, Bharat Venkitesh, Dwarak Talupuru, Hangyu Lin, David Cairuz, Phil Blunsom, and Acyr Locatelli. Rope to nope and back again: A new hybrid attention strategy. arXiv preprint arXiv:2501.18795, 2025. [50] Mingyu Yang, Mehdi Rezagholizadeh, Guihong Li, Vikram Appia, and Emad Barsoum. Zebra-llama: Towards extremely efficient hybrid models. arXiv preprint arXiv:2505.17272, 2025. [51] Songlin Yang and Yu Zhang. Fla: A triton-based library for hardware-efficient implementations of linear attention mechanism, January 2024. URL https://github.com/fla-org/flash-linear-attention. [52] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06464, 2024. [53] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830, 2019.
13
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
A
Appendix
A.1
More Experimental Details
A.1.1
Details of Model Configurations
We implement our upcycling recipe starting from three different base models: Llama-3.2-1B, Llama-3.2-3B, and Qwen3-1.7B. Notably, Qwen3 utilizes normalization after the Query and Key projection layers, which we maintain when performing attention layer conversion. Table 7 summarizes the model configurations and hyperparameters used in our experiments.
A.1.2
Enhanced-ILD hyperparameters.
MLA/Mamba2/GDN Enhanced-ILD runs share the same configuration: 1 epoch on 20% of the training SFT dataset [46], with context length 2048, learning rate 2×10−4 with cosine decay, warmup ratio 0.01, and bfloat16 mixed precision. Training is distributed across 8 AMD MI300X GPUs using FSDP with full sharding.
A.1.3
SFT hyperparameters
We train for 1 epoch at context length 8K/64K, learning rate as reported in Table 7 with cosine schedule, warmup ratio 0.01, and the full dataset (data ratio=1.0). YaRN-based position scaling extends the effective context from the original 2048 to 8192 tokens (scaling factor 4.0) or from the original 2048 to 65536 (scaling factor 32.0). Training uses FSDP across 8 AMD MI300X GPUs. Our Models
Base Model
MLA Layer Indices
HyLo-Llama-4MLA12M2 HyLo-Llama-4MLA12GDN HyLo-Llama-8MLA8M2 HyLo-Llama-8MLA8GDN HyLo-Llama-6MLA22M2 HyLo-Llama-6MLA22GDN
Llama-3.2-1B Llama-3.2-1B Llama-3.2-1B Llama-3.2-1B Llama-3.2-3B Llama-3.2-3B
HyLo-Llama-14MLA14M2
Llama-3.2-3B
[1,5,10,14] [1,5,10,14] [0,2,4,6,8,10,12,14] [0,2,4,6,8,10,12,14] [0,5,10,16,21,26] [0,5,10,16,21,26] [0,2,4,6,8,10,12,14, 16,18,20,22,24,26] [0,2,4,6,8,10,12,14, 16,18,20,22,24,26] [1,5,9,13,17,21,25] [1,5,9,13,17,21,25] [0,2,4,6,8,10,12,14, 16,18,20,22,24,26] [0,2,4,6,8,10,12,14, 16,18,20,22,24,26]
HyLo-Llama-14MLA14GDN
Llama-3.2-3B
HyLo-Qwen-7MLA21M2 HyLo-Qwen-7MLA21GDN
Qwen3-1.7B Qwen3-1.7B
HyLo-Qwen-14MLA14M2
Qwen3-1.7B
HyLo-Qwen-14MLA14GDN
Qwen3-1.7B
# Act. Params.
Head
Layer
Hidden
lr
batch size (8k/64k)
1.5B 1.7B 1.5B 1.6B 3.8B 4.3B
32 32 32 32 24 24
16 16 16 16 28 28
2048 2048 2048 2048 3072 3072
6.0 × 10−5 6.0 × 10−5 6.0 × 10−5 6.0 × 10−5 4.0 × 10−5 4.0 × 10−5
32/8 32/8 32/8 32/8 16/8 16/8
3.7B
24
28
3072
4.0 × 10−5
16/8
−5
4.0B
24
28
3072
4.0 × 10
16/8
2.1B 2.3B
16 16
28 28
2048 2048
6.0 × 10−5 6.0 × 10−5
32/8 16/8
2.1B
16
28
2048
6.0 × 10−5
32/8
2.2B
16
28
2048
6.0 × 10−5
16/8
Table 7: Model configurations and hyperparameters for our experiments.
A.2
MLA Layer Architecture and SVD-Based Initialization
The Multi-head Latent Attention (MLA) layers in HyLo follow the DeepSeek-V3 design [30], which compresses the key-value cache through low-rank latent projections. We initialize MLA layers from pretrained Transformer attention weights using SVD-based decomposition following the methodology outlined in X-EcoMLA [25].
A.2.1
MLA Architecture
Given input xt ∈ Rd , the MLA module computes queries, keys, and values through two low-rank bottlenecks: 14
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Query path. QA cQ xt ∈ Rrq , t =W nope
H×dqk qnope = WQB Norm(cQ , t t )∈R
qrope = WQR Norm(cQ t t )∈R
H×drope qk
(5)
,
where rq is the query latent rank, and qrope receives RoPE. t Key-value path. cKV = WKV A xt ∈ Rrkv , t rope
krope = WKR xt ∈ Rdqk , t
(6)
nope
knope = WKB Norm(cKV ) ∈ RHkv ×dqk , t t vt = WV B Norm(cKV ) ∈ RHkv ×dv , t
where rkv is the joint key-value latent rank. The KV cache only stores the compressed latent cKV ∈ Rrkv and the rope t rope rope drope key kt ∈ R qk , reducing the per-token cache from 2Hkv dh to rkv + dqk . Attention and output. The full query and key are assembled as qt = [qnope ; RoPE(qrope t t )] and kt = nope rope [kt ; RoPE(kt )], then standard scaled dot-product attention is applied: ot = WO Attn(qt , kt , vt ). A.2.2
(7)
SVD-Based Initialization from Transformer Weights
To initialize the low-rank MLA projections from a pretrained Transformer model, we decompose the teacher’s full-rank attention weights via truncated SVD. Query initialization.
Let WQ ∈ R(H·dh )×d be the teacher’s query projection. We compute its SVD: ⊤ WQ = UQ ΣQ VQ ,
(8)
and initialize the MLA down/up projections as: WQA ← ΣQ [: rq ] VQ [: rq , :]⊤ ∈ Rrq ×d ,
(9)
rope (H·dqk )×rq WQB ← Select(UQ [:, : rq ], dnope , qk , dqk ) ∈ R
rope where Select(·) reshapes UQ into per-head blocks and retains only the first dnope qk and last dqk dimensions from each head’s dh -dimensional slice, discarding the middle dimensions that are not used in MLA.
Joint key-value initialization. Key-value initialization is complicated by MLA’s decoupled RoPE design. When the teacher uses GQA (Hkv < H), we first expand WK and WV to H heads by replicating each KV group H/Hkv times, then apply truncated SVD to the concatenated matrix: ⊤ [WK , WV ] = UKV ΣKV VKV .
(10)
We set WKV A = UKV [:, : rkv ], With dv = dh , we split W
KV B
WKV B = ΣKV [: rkv ] VKV [: rkv , :]⊤ . K
(11)
V
into key and value parts (same column order as [W , W ]):
WV B = WKV B [:, Hkv dh :], W̄KB = reshape(WKV B [:, : Hkv dh ]) ∈ Rrkv ×Hkv ×dh , W
KB
= reshape(W̄
KB
15
[:, :, : dqk ]).
(12)
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Figure 6: Overview of MLA initialization from a pretrained Transformer attention block. Finally, because all MLA heads share the same RoPE key embedding, we initialize WKR from the head-averaged key K projection Wavg : K WKR = Wavg [:, −dr :]. (13) Output projection.
The output projection is truncated from the teacher: WO ← WO [:, : H · dv ] ∈ Rd×(H·dv ) .
MLP and layer norms.
A.3
(14)
All MLP weights and RMSNorm parameters are copied directly from the teacher.
GDN Layer Architecture
In the GDN-based HyLo hybrid architecture, each non-attention decoder layer replaces the standard attention module with a Gated DeltaNet (GDN) [52] mixer while preserving the SwiGLU MLP and RMSNorm sub-layers from the original transformer block. Concretely, a GDN decoder layer consists of: h′ = h + GDN RMSNorm(h) , (15) h′′ = h′ + MLP RMSNorm(h′ ) , where GDN(·) is the Gated DeltaNet module from FLA [51]. GDN mixer parameterization. With the gating mechanism enabled (use gate=True), the GDN mixer allocates parameters as follows. Let d denote the model hidden size. The key dimension is dk = ⌊0.75 · d⌋, distributed over H heads each of dimension dh = dk /H. The value dimension is dv = α · dk with expansion ratio α = 2. The projections are: 16
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Parameter Q
K
W ,W WV , WG , WO Wα , Wβ Alog , ∆bias Short conv (Q, K, V )
Shape
Count
1536 × 2048 3072 × 2048 6 × 2048 6 kernel 4
2 × 3.15M 3 × 6.29M 2 × 12.3K 12 ∼ 31K ∼ 25.2M
GDN mixer total
Table 8: GDN parameter dimensions for Llama-3.2-1B (d = 2048, H = 6, dh = 256, α = 2). • WQ , WK ∈ Rdk ×d (0.75 d2 parameters each), • WV , WG , WO ∈ Rdv ×d (1.5 d2 parameters each), • Wα , Wβ ∈ RH×d (decay and beta projections), • Alog ∈ RH , ∆bias ∈ RH (learnable decay and timestep biases), yielding approximately 6 d2 parameters per layer. Each of Q, K, and V is processed through a short 1-D convolution (kernel size 4) with SiLU activation before the recurrence. Gated Delta Rule Recurrence. The GDN mixer maintains a per-head state matrix St ∈ Rdk ×dv that is updated at every timestep t via the gated delta rule: S̃t = egt · St−1 ,
(16)
vt′ = vt − S̃⊤t kt ,
(17) ′ ⊤
St = S̃t + kt βt · vt
,
(18)
1 ot = √ S⊤t qt , dk ot = RMSNorm qt St , WG xt ,
(20)
yt = WO ot ,
(21)
(19)
where gt ∈ (−∞, 0) is the per-head forget gate (Eq. 16), βt ∈ (0, 1) is the per-head write strength, qt , kt ∈ Rdk are queries and keys, vt ∈ Rdv are values, and dk = ⌊0.75 d⌋, dv = 2 dk . Eq. (16) applies an exponential decay to the state, controlled by gt = − exp(Alog ) · softplus(Wα xt + ∆bias ). Eq. (17) is the delta correction: it retrieves the value currently associated with key kt and subtracts it from the new value vt , preventing superposition interference. Eq. (18) writes the corrected value back into the state, scaled by βt = σ(Wβ xt ). Eq. (20) reads the output by querying the updated state with qt . During training, the sequential recurrence is computed efficiently using the chunked kernel of Yang et al. [52], which partitions the sequence into chunks of size C=64. Within each chunk, the delta rule corrections are batched into a single matrix operation via the WY representation; across chunks, the state S is propagated sequentially over T /C steps instead of T , yielding linear-time complexity with high GPU utilization.
A.4
Memory-Efficient Long-Context Knowledge Distillation
Extending knowledge distillation from 2K to 64K context lengths introduces severe memory pressure. The dominant bottleneck is the logit tensor: for sequence length T and vocabulary size V , the standard KL divergence loss requires materializing both student and teacher logit matrices of shape (T, V ) simultaneously. At T =65,536 and V =128,256 (Llama-3), each logit tensor consumes approximately 16 GB in bfloat16—making naive distillation infeasible even on 80 GB GPUs. We address this through a progression of increasingly aggressive memory optimizations, summarized in Table 9 with training memory reported for some of the optimizations in Table 1. 17
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
A.4.1
Fused Linear Cross-Entropy
The standard cross-entropy computation first projects hidden states through the LM head Wlm ∈ RV ×d to produce ⊤ the full logit matrix Z = HWlm ∈ RT ×V , then applies the softmax and loss in a separate step. This requires O(T V ) memory just for the logit materialization. A.4.2
Chunked KL Divergence
The KL divergence DKL (ps ∥pt ) between student and teacher logits normally requires both (T, V ) log-softmax tensors in memory. For long contexts, we chunk along the sequence dimension: 1 DKL = T
⌈T /C⌉−1 min((i+1)C,T )−1 (s) (t) DKL softmax(zj )∥softmax(zj ) ,
X
X
i=0
j=iC
(22)
with C=4,096. Each chunk allocates only a (C, V ) softmax slice, reducing peak memory from 2 × T × V to 2 × C × V . Intermediate tensors are explicitly freed between chunks. A.4.3
Triton-Fused KL Divergence
For further efficiency, we leverage a custom Triton kernel from FLA [51] that computes DKL entirely within a single fused kernel using online softmax [32]. The kernel tiles over the vocabulary dimension with block size BV and maintains running log-sum-exp accumulators, so the full (T, V ) softmax matrices are never materialized: (j)
DKL =
⌈V /BV ⌉−1 h
X
i (j) accum log p(j) [bB : (b+1)B ], p [bB : (b+1)B ] , V V V V t s
(23)
b=0
where each token j is processed by one Triton program instance, and the gradient ∂DKL /∂z(s) is written in-place during the forward pass (overwriting the student logit buffer), eliminating the need to save activations for backward. A.4.4
Fused Hidden-State KL (Logit-Free Distillation)
At the longest contexts (64K tokens), even chunked approaches are bottlenecked by the need to run the teacher’s LM head. We use a logit-free distillation path: the teacher forward pass skips the LM head entirely, returning only the final hidden states H(t) ∈ RT ×dt . The FLA FusedKLDivLoss then computes KL directly from hidden states and LM head weight matrices: (s)⊤ (t)⊤ L = DKL softmax H(s) Wlm softmax H(t) Wlm , (24) where the softmax and KL are computed in a tiled fashion inside the Triton kernel without materializing either logit matrix. This eliminates 2 × T × V elements from GPU memory (approximately 32 GB at T =64K). The teacher’s LM (t) head weight Wlm is accessed via FSDP’s summon full params to avoid duplicating sharded parameters. A.4.5
Teacher Memory Management
To accommodate the teacher model (Llama-3.1-8B) alongside the student HyLo models during long-context distillation, we employ several additional strategies: • Frozen teacher under torch.no grad: The teacher runs in evaluation mode with gradient computation disabled, eliminating all optimizer states, gradient tensors, and backward graph storage for teacher parameters. • FSDP full sharding for both models: Both student and teacher are wrapped with FSDP using FULL SHARD strategy, distributing model parameters across all 8 GPUs. Only one shard per GPU is materialized at a time during forward/backward passes. 18
Long-Context Aware Upcycling: A New Frontier for Hybrid LLM Scaling
Table 9: Memory optimization techniques and their deployment across context lengths. Each technique targets a specific memory bottleneck in the knowledge distillation pipeline. Technique
Memory Saved
Used at
Liger Fused Linear CE Chunked KL Divergence Triton Fused KL Fused Hidden-State KL FSDP Full Sharding Frozen teacher (no grad) bf16 mixed precision
Student logits (T ×V ) Softmax tensors 2(T ×V ) Softmax + grad 3(T ×V ) Both logit matrices 2(T ×V ) Model params ÷NGPU Teacher grads + optimizer 2× vs. fp32
8K–32K 64K 128K 64K All All All
• bfloat16 mixed precision: All activations and parameters are stored in bfloat16, halving memory relative to float32. • Batch size reduction: At 64K+ contexts, per-device batch size is reduced to 1 (from 4 at 8K), trading throughput for memory headroom. • Activation checkpointing (optional): Gradient checkpointing is supported for the decoder stack and can be enabled when activation memory dominates; however, the sub-quadratic memory of linear recurrence layers (Mamba-2/GDN) and the fused loss kernels typically suffice without it.
A.5
LLM Usage
The authors of this paper used AI tools for polishing text within this paper. The authors take full responsibility for the content within this paper.
19