KITE：用于高效智能体LLM扩展的KV不变Transformer扩展方法

# KV-Invariant Transformer Expansionfor Efficient Agentic LLM Scaling
Source: [https://arxiv.org/html/2609.27294](https://arxiv.org/html/2609.27294)
Zhiheng Hu Yixun Wei Jian Zhou Yizhuang Zhou Ji Li Xing ChenYang Li Bojun Wang Yibo Zhu Xiangyu Zhang Daxin JiangStepFun![[Uncaptioned image]](https://arxiv.org/html/2609.27294v1/stepfun_logo.png)

###### Abstract

Scaling a language model is not only a question of final quality: the architectural choice determines how much computation is spent during training, prompt processing, and autoregressive decoding to achieve certain model quality\. An ideal model architecture should lower all above computation costs to facilitate scaling to a larger model, while ensure the larger model indeed outperforms smaller baselines\. We introduce*KV\-Invariant Transformer Expansion \(KITE\)*, a scaling paradigm that achieves this goal\. It trains the model from a smaller size to a larger size \(i\.e\.,saving training costs via upcycling\), while places newly added parameters in regions that do not affect attention KV\. Consequently, during inference, prefilling KV only relies on the smaller part of the model, so the inference costs are saved\. As a concrete instantiation, we present Step Scale Transformer \(SST\), a two\-tower decoder in which one tower produces KV and the other reads them\. At comparable cumulative training compute, SST, a 67B MoE model with 2\.15B active body parameters per decode token, achieves lower training loss than 47B and 63B MoE Transformers with 1\.48B and 2\.02B active body parameters, respectively, while reducing estimated inference cost by 6\.7% and 31\.6%\.

Figure 1:KITE: training route and observed trade\-offs\. Both panels use the same 47B transformer trained from scratch as the reference; lower is inward\. Active counts exclude embedding/head\. The inference proxy is0\.75​p\+0\.25​d0\.75p\+0\.25d, whereppandddare bulk\-prefill and decode active\-body counts normalized to the reference\.## 1Introduction

Large language models \(LLMs\) exhibit strong scaling behavior: increasing model scale continues to bring substantial gains in model capability\[[14](https://arxiv.org/html/2609.27294#bib.bib1),[13](https://arxiv.org/html/2609.27294#bib.bib2)\]\. This property is both one of the most important blessings of LLMs and a curse, as larger models require enormous compute budgets for both training and inference\.

Consequently, one of the central problems in LLM research is how to scale models at lower compute cost\. This goal is not only about reducing expenditure; it enables a larger model to be trained or served under the same compute budget, yielding a stronger model\. MoE feed\-forward layers provide a successful example: replacing a dense FFN with sparse expert routing substantially reduces training and inference computation at a fixed parameter count, or, equivalently, allows model capacity to grow at a fixed compute budget\[[19](https://arxiv.org/html/2609.27294#bib.bib13),[9](https://arxiv.org/html/2609.27294#bib.bib5)\]\.

What is easy to miss in this success is that MoE reduces training and inference computation simultaneously\. A saving on only one side may not improve the overall trade\-off\. For example, progressive upcycling trains a small model first and grows it into the target model partway through training, so much of the token budget runs on a cheaper graph\[[6](https://arxiv.org/html/2609.27294#bib.bib3),[28](https://arxiv.org/html/2609.27294#bib.bib4)\]\. However, upcycled models are generally perceived to underperform same\-size models trained from scratch, so the training saving is paid back at inference, where one serves the full target model yet obtains less than the full target quality\. KV reuse pays on the other side\. By letting deeper layers reuse the keys and values computed by shallower layers, YOCO family keeps prefill affordable as a model grows\[[22](https://arxiv.org/html/2609.27294#bib.bib7),[25](https://arxiv.org/html/2609.27294#bib.bib9),[8](https://arxiv.org/html/2609.27294#bib.bib10)\]; Yet removing KV\-producing layers can sacrifice model quality; recovering the same quality may require a larger model or a longer training, which is effectively a training\-cost loss\. In both cases the gain on one axis is offset by a regression on the other\. We therefore treat the two costs jointly: we pursue a scaling paradigm whose model quality, training cost, and inference cost all move in its favor\.

This paper proposes*KV\-Invariant Transformer Expansion \(KITE\)*, a model\-scaling paradigm\. Its core idea is that when the model is scaled, it is divided into two regions: the region that determines the key–value \(KV\) tensors, and the region that does not affect them\. Scaling is carried out by expanding capacity in the region that does not affect the KV, which yields a good ratio of model intelligence to both training and inference cost\. That is because the added capacity appears partway through training rather than from token zero, so the quality it buys comes at a favorable training cost\. Meanwhile, this scaling does not affect the KV computation, so the prefill cost at inference remains unchanged\.

KITE leverages the fact that prompt processing is significant in agentic workloads, where repeated model calls process accumulated context and new tool outputs\[[27](https://arxiv.org/html/2609.27294#bib.bib8),[26](https://arxiv.org/html/2609.27294#bib.bib18)\]\. Even with prompt caching, substantial uncached input remains\. Public OpenRouter traffic illustrates this input\-heavy demand: across the six models in Table[1](https://arxiv.org/html/2609.27294#S1.T1), uncached input exceeds output in token volume and accounts for most estimated uncached\-input and output charges\. Consequently, by saving on the prefill costs, KITE, like YOCO, produces models that are overall much cheaper to serve than classic transformers\.

Table 1:Input demand in selected OpenRouter traffic\. Input includes cached tokens; output includes reasoning tokens\. Charge share is uncached\-input charges divided by uncached\-input plus output charges, excluding cache\-read charges\. Estimates aggregate available standard and Batch activity within August 22–September 20, 2026 \(UTC\), using September 21 base prices\. Details appear in Appendix[A](https://arxiv.org/html/2609.27294#A1)\.To instantiate KITE, we design a straightforward architecture:*Step Scale Transformer*\(SST\)\. SST contains two same\-depth towers that execute sequentially\. The Prefiller produces layer\-wise KV, which the Decoder reuses to predict the next token\. During generation, bulk prefill uses only the first tower, while the prompt\-boundary prediction and decode use both\. The source model is trained first; after expansion, both towers continue training jointly\. Section[2](https://arxiv.org/html/2609.27294#S2)provides the architecture and execution details\.

We spend the saved compute on scaling the model up moderately, just as MoE does\. Compared with a Transformer baseline trained from scratch, this yields a net gain in the three\-way trade\-off among training cost, inference cost, and model intelligence\. Figure[1](https://arxiv.org/html/2609.27294#S0.F1)illustrates these trade\-offs\. KITE can also be used together with earlier successful designs such as MoE and hybrid attention; they are orthogonal and do not interfere with one another\. Although SST is probably not the best KITE design, it is straightforward and suffices to demonstrate the positive effect of the paradigm\. We are continuing to explore this paradigm, and we anticipate that more and more related designs can emerge\.

To make the comparison clear, we build a scaling ladder across multiple model scaling paradigms\. The tokenizer, data recipe, training protocol, and evaluation sets are kept identical, and we trace the cost–quality relationship of each scaling paradigm\. A better scaling paradigm is one that reaches a lower training loss under a comparable training\-FLOPs budget while also incurring a lower inference cost\. The experiments show that the SST instantiation of KITE, expanded to 67B parameters, reaches an EMA\-200 training loss of 1\.5900, 0\.0106 lower than 47B classic \(1\.6006\) and 0\.0021 lower than 63B classic \(1\.5921\)\. At the final training endpoints, SST and 63B classic both use approximately 100% of the 47B baseline’s theoretical cumulative training FLOPs, including both SST training stages\. With illustrative cost weights of 75% for prefill and 25% for decode, SST’s analytical inference\-cost proxy is 6\.7% lower than 47B classic and 31\.6% lower than 63B classic \(Figure[1](https://arxiv.org/html/2609.27294#S0.F1)\)\. On downstream benchmarks, SST scores higher than both baselines on OpenBookQA, MMLU, GSM8K, MATH, HumanEval, MBPP, and BBH \(Table[3](https://arxiv.org/html/2609.27294#S4.T3)\)\.

## 2KITE: KV\-Invariant Transformer Expansion

We illustrate KITE through SST, a concrete instantiation that adds KV\-reading capacity to a smaller model and continues training both parts jointly \(Figure[2](https://arxiv.org/html/2609.27294#S2.F2)\)\.

Stage 1Train the PrefillerStage 2Jointly train both towersPrefillerTransformer block 1Transformer block 2Transformer blockLL⋮\\vdotsToken embeddingOutput norm \+ LM headNext\-token predictionExpandInitialize Decoderfrom source weightsretainedadded at expansionPrefillerDecoderTransformer block 1Transformer block 2Transformer blockLLTransformer block 1Transformer block 2Transformer blockLL⋮\\vdots⋮\\vdotsToken embeddingEntry bridgeRMSNorm \(no affine\)Output norm \+ LM headNext\-token predictionFinal Prefiller stateEmbeddingK1,V1K\_\{1\},V\_\{1\}K2,V2K\_\{2\},V\_\{2\}KL,VLK\_\{L\},V\_\{L\}Layer\-wiseKV reuseShared KVToken\-local bridgeNew Decoder

Figure 2:KITE instantiated with SST\. Stage 1 trains the Prefiller; expansion adds a Decoder initialized from the source weights, and Stage 2 continues training both towers jointly\. Within Stage 2, the Prefiller executes in full before the Decoder\. Each Decoder block uses aligned Prefiller KV\. A token\-local entry bridge combines the separately normalized embedding and final Prefiller state\. The final Decoder state is normalized and passed to the shared output norm and LM head for prediction\.### 2\.1The design principle

Let a causal model have a PrefillerPPand a DecoderDD:

\(Mt,zt\)=P\(x1:t;θP\),ℓt=D\(xt,zt,Mt;θD\),\(M\_\{t\},z\_\{t\}\)=P\(x\_\{1:t\};\\theta\_\{P\}\),\\qquad\\ell\_\{t\}=D\(x\_\{t\},z\_\{t\},M\_\{t\};\\theta\_\{D\}\),\(1\)whereMtM\_\{t\}is reusable attention memory,ztz\_\{t\}denotes optional Prefiller features at positiontt, andℓt\\ell\_\{t\}gives next\-token logits\. After an initial small\-model training stage, expansion adds or enlargesDDusing the intermediate checkpoint, retaining the Prefiller graph and memory interface\. Both parts then continue training\.

The invariant is structural: neither weights nor KV values are frozen\. The Decoder must not produce memory needed by future positions\. To omit historical Decoder computation during generation, its retained positions must also be independent of omitted Decoder positions given the Prefiller outputs\. These restrictions let prediction capacity grow without increasing*bulk prefill*, the prompt\-wide KV construction performed byPP\. The final\-position readout for the first output token and the Decoder’s work during decode remain necessary\.

### 2\.2Step Scale Transformer

SST instantiates this principle with two Transformer towers\[[24](https://arxiv.org/html/2609.27294#bib.bib12)\]: the Prefiller \(superscriptPP\) and the Decoder \(DD\), each with 18 layers in our experiments\. The Prefiller executes first and produces layer\-wise KV\. Each Decoder block computes its own query and reads KV from the corresponding Prefiller layer:

ql,tD=QlD​\(hl−1,tD\),al,tD=Attn⁡\(ql,tD,KlP​\[ℐl​\(t\)\],VlP​\[ℐl​\(t\)\]\)\.q^\{D\}\_\{l,t\}=Q\_\{l\}^\{D\}\(h^\{D\}\_\{l\-1,t\}\),\\qquad a^\{D\}\_\{l,t\}=\\operatorname\{Attn\}\\\!\\left\(q^\{D\}\_\{l,t\},K\_\{l\}^\{P\}\[\\mathcal\{I\}\_\{l\}\(t\)\],V\_\{l\}^\{P\}\[\\mathcal\{I\}\_\{l\}\(t\)\]\\right\)\.\(2\)
Here,QlDQ\_\{l\}^\{D\}includes query projection and normalization, andℐl​\(t\)\\mathcal\{I\}\_\{l\}\(t\)specifies the positions allowed by the original causal full/sliding attention mask\. Both towers retain the source model’s position IDs\. The Decoder has independent block parameters; its residual and FFN/MoE operations are token\-local\.

In the expanded SST model shown on the right of Figure[2](https://arxiv.org/html/2609.27294#S2.F2), a token\-local entry bridge initializes the Decoder from the embedding and final Prefiller state\. Prediction uses the Decoder’s final state\. Gradients flow to the Prefiller through both the entry bridge and the reused KV\. Given the Prefiller outputs, each Decoder position is independent of earlier Decoder states\.

### 2\.3Training and inference

#### Training\.

After expansion, SST jointly trains both towers\. The Prefiller processes the full sequence, and the Decoder computes every supervised position; gradients pass through the reused KV into the Prefiller\. The initial small\-model training stage precedes the addition of the Decoder\.

#### Prefill\.

The Prefiller processes the complete promptx1:nx\_\{1:n\}and produces the layer\-wise KV cache\. The Decoder needs to compute only the final prompt positionnn, attending to the Prefiller’s KV, to predict the first output token\. Earlier Decoder positions can be omitted because its other operations are token\-local\.

#### Decode\.

Each new token passes through both towers\. The Prefiller extends the KV cache, and the Decoder reads this KV to form the next\-token prediction\.

### 2\.4Decoder input and output

For the shared embedding/stemete\_\{t\}and final Prefiller statehL,tPh^\{P\}\_\{L,t\}, SST uses the entry bridge

h0,tD=RMSϵ⁡\(et\)\+RMSϵ⁡\(hL,tP\),RMSϵ⁡\(v\)=vmean⁡\(v2\)\+ϵ\.h^\{D\}\_\{0,t\}=\\operatorname\{RMS\}\_\{\\epsilon\}\(e\_\{t\}\)\+\\operatorname\{RMS\}\_\{\\epsilon\}\(h^\{P\}\_\{L,t\}\),\\qquad\\operatorname\{RMS\}\_\{\\epsilon\}\(v\)=\\frac\{v\}\{\\sqrt\{\\operatorname\{mean\}\(v^\{2\}\)\+\\epsilon\}\}\.\(3\)The Decoder’s final state produces the next\-token logits:

ℓt=Wout​Normout​\(hL,tD\)\.\\ell\_\{t\}=W\_\{\\mathrm\{out\}\}\\,\\mathrm\{Norm\}\_\{\\mathrm\{out\}\}\(h^\{D\}\_\{L,t\}\)\.\(4\)Here,Normout\\mathrm\{Norm\}\_\{\\mathrm\{out\}\}denotes the complete output normalization; its implementation is specified in Appendix[D](https://arxiv.org/html/2609.27294#A4)\. The entry\-bridge RMS operations have no trainable parameters and useϵ=10−5\\epsilon=10^\{\-5\}\. The towers share one embedding, output normalization, and output head; embedding and head weights are untied\.

## 3Model construction and training

### 3\.1From the source to the expanded model

SST expands a 33\.819B\-parameter source into a 66\.959B model while retaining the source tower as the Prefiller \(Table[2](https://arxiv.org/html/2609.27294#S3.T2)\)\. The evaluated route contains 167\.98B source\-stage tokens and 222\.55B continuation tokens \(Figure[3](https://arxiv.org/html/2609.27294#S3.F3)\)\. Both stages contribute to the reported training budget\.

The 47B and 63B classic\-trained baselines are the next two scales above the 33\.8B source in the same model family’s scaling ladder, obtained by increasing depth and width while retaining the MoE design and attention pattern\. The 47B model provides an intermediate\-scale reference, and the 63B model provides a reference close to SST’s 67B parameter count\.

Table 2:Model configurations\.Each tower starts with two dense FFN layers\. SSSF is three sliding layers followed by one full\-attention layer, with an SF remainder for 18/22 layers\. Counts describe the effective architecture; Appendix[B](https://arxiv.org/html/2609.27294#A2)gives the geometry and counting convention\.

Source training167\.98B tokensJoint Prefiller/Decoder training222\.55B tokens0167\.98B390\.54BConvert checkpointCopy weights into both towersEvaluated model31,037 total updatesCumulative training tokens; source and continuation are both included\.Figure 3:Training route to the evaluated SST checkpoint: 13,350 source steps followed by 17,687 continuation steps\.
### 3\.2Expansion and continued training

We expand the source checkpoint into SST by copying each source layer into the corresponding layers of the Prefiller and the Decoder\. The two towers share a single copy of the embedding, output normalization, and output head\.

After expansion, both towers are trained jointly in BF16 using the source model’s Muon\[[17](https://arxiv.org/html/2609.27294#bib.bib21)\]and Adam\[[15](https://arxiv.org/html/2609.27294#bib.bib20)\]parameter groups\. We retain the optimizer states for the Prefiller and the shared parameters and initialize fresh states for the Decoder\. All groups continue on the source cosine learning\-rate schedule, with the remaining decay extended over the continuation stage and no separate warmup for the Decoder\. Appendix[D](https://arxiv.org/html/2609.27294#A4)gives the detailed settings\.

### 3\.3Cumulative training compute

We compare the theoretical FLOPs of the effective architectures\. For source and continuation token countsDs,DcD\_\{s\},D\_\{c\}and forward FLOPs/tokenfs,fcf\_\{s\},f\_\{c\}, we use

Ctrain=3​\(Ds​fs\+Dc​fc\),C\_\{\\mathrm\{train\}\}=3\(D\_\{s\}f\_\{s\}\+D\_\{c\}f\_\{c\}\),\(5\)where the factor three approximates forward\-plus\-backward computation\. The full cost of both training stages is included\. The evaluated 47B classic checkpoint definesBref≈5\.16836×1021B\_\{\\mathrm\{ref\}\}\\approx 5\.16836\\times 10^\{21\}FLOPs\. The final SST and 63B endpoints both use approximately 100% of this budget\. Appendix[C](https://arxiv.org/html/2609.27294#A3)gives the theoretical accounting and training budgets\.

## 4Results

### 4\.1Training loss

SST’s loss rises at conversion and then decreases through continuation \(Figure[4](https://arxiv.org/html/2609.27294#S4.F4)\)\. At the final checkpoints, its training EMA\-200 is 1\.5900 after 390\.54B tokens, versus 1\.6006 for 47B classic after 443\.16B tokens and 1\.5921 for 63B classic after 334\.62B tokens\. SST therefore reaches lower training loss than both references at comparable cumulative training compute\. Training budgets and loss reporting are given in Appendix[C](https://arxiv.org/html/2609.27294#A3)\.

Figure 4:Training LM loss \(EMA\-200\) against tokens and theoretical cumulative compute\. Conversion occurs at 167\.98B tokens \(0\.301​Bref0\.301B\_\{\\mathrm\{ref\}\}\)\.
### 4\.2Task\-level evaluation

Table[3](https://arxiv.org/html/2609.27294#S4.T3)reports scores on seven downstream tasks \(OpenBookQA\[[18](https://arxiv.org/html/2609.27294#bib.bib22)\], MMLU\[[11](https://arxiv.org/html/2609.27294#bib.bib23)\], GSM8K\[[7](https://arxiv.org/html/2609.27294#bib.bib24)\], MATH\[[12](https://arxiv.org/html/2609.27294#bib.bib25)\], HumanEval\[[5](https://arxiv.org/html/2609.27294#bib.bib26)\], MBPP\[[2](https://arxiv.org/html/2609.27294#bib.bib27)\], and BBH\[[23](https://arxiv.org/html/2609.27294#bib.bib28)\]\) at the final checkpoints of all three models\. These are the same checkpoints used for the training\-loss and cumulative\-compute comparisons\. We also evaluate held\-out token NLL on ARXIV using the same 1,024 frozen packed\-4K samples for all three models\.

Trained with an equal total of FLOPs, SST performs the best across the board\. It shows that SST’s lowest training loss indeed translates to the best downstream performance\. SST also has the lowest estimated inference cost in the input\-heavy regime analyzed in Section[4\.3](https://arxiv.org/html/2609.27294#S4.SS3)\.

Table 3:Scores \(%\) on seven downstream tasks and held\-out ARXIV token NLL \(lower is better\), evaluated at comparable cumulative theoretical training FLOPs, including both training stages for SST\. OpenBookQA combines the validation and test splits; MMLU uses the test split, and MATH uses the full test set\.
### 4\.3Inference cost across prefill–decode mixes

Figure[5](https://arxiv.org/html/2609.27294#S4.F5)compares estimated inference cost across prefill:decode \(P:D\) cost mixes, using the proxy in Appendix[A](https://arxiv.org/html/2609.27294#A1)and active\-body counts in Appendix[B](https://arxiv.org/html/2609.27294#A2)\. All models use the same P:D weights, with costs normalized to 47B classic\. SST breaks even with 63B at 13\.4:86\.6 and with 47B at 65\.5:34\.5, becoming cheaper as the prefill share increases\. At 75:25, its estimated cost is 6\.7% lower than 47B and 31\.6% lower than 63B\.

Figure 5:Estimated inference cost across P:D cost mixes, normalized to 47B classic\. Shading spans the six\-model OpenRouter reference range; open circles mark break\-even points and the filled circle marks 75:25\.The six OpenRouter models in Table[1](https://arxiv.org/html/2609.27294#S1.T1)have uncached\-input bill shares of 67\.4–76\.7% \(Appendix[A](https://arxiv.org/html/2609.27294#A1)\)\. Using these shares as reference prefill weights, SST is cheaper than both baselines throughout the highlighted range: 1\.3–7\.9% below 47B and 27\.7–32\.5% below 63B\. These are analytical cost estimates, not measured serving speedups\.

## 5Discussion

SST provides one concrete realization of KITE, rather than a unique architectural prescription\. Its aligned layer\-wise KV reuse, entry bridge, and Decoder\-only readout are specific design choices\. The central requirement is that the added capacity does not extend the KV\-producing computation or introduce dependencies that require historical Decoder states\. Figure[6](https://arxiv.org/html/2609.27294#S5.F6)illustrates several connectivity choices within this framework\.

\(a\) Connectivity choices\(b\) KV assignmentThree\-layer examples; choose one mappingPrefillerDecoderProduces KVReads Prefiller KVTransformer block 1Transformer blockllTransformer blockLLTransformer block 1Transformer blockllTransformer blockLL⋮\\vdots⋮\\vdots⋮\\vdots⋮\\vdotsToken embeddingInput combinationReadout \+ LM headNext\-token predictionEntry bridgeIntermediate bridgeOutput fusionK1,V1K\_\{1\},V\_\{1\}Kl,VlK\_\{l\},V\_\{l\}KL,VLK\_\{L\},V\_\{L\}KV access retained;assignment can changePD111122223333PD111122223333PD111122223333BackboneKV accessOptional hidden\-state pathBridge / readout operators may varyBoundary: future KV production must not depend on historical Decoder states\.

Figure 6:Connectivity choices within KITE\. \(a\) An SST\-style backbone with optional hidden\-state connections at the Decoder entry, intermediate layers, and output readout\. Dashed paths indicate design choices, not connections that must be enabled together\. \(b\) Three illustrative KV assignments: aligned layer\-wise reuse, reversed pairing, and reuse of one Prefiller KV source across multiple Decoder layers\. The KV\-producing computation remains on the Prefiller side; the illustrated choices are not assumed to have equal cost or quality\.The assignment between Prefiller KV and Decoder layers need not be one\-to\-one or order\-preserving\. For example, Decoder layers could read Prefiller KV in reversed layer order, or multiple Decoder layers could reuse a single KV source\. Hidden\-state connections offer another degree of freedom: Prefiller representations may enter the Decoder at its input or at intermediate layers, and may also contribute to the final readout\. The bridge operators, as well as the depth and width of the two components, can likewise vary\. These choices must preserve causal access and compatible interfaces, but need not have the same computational cost or model quality\.

Our experiments explore only a limited part of this design space\. The results reported in this paper establish the performance of the selected SST configuration, rather than identify an optimal connectivity pattern\. Further exploration can examine how KV assignment, hidden\-state connections, and capacity allocation affect the balance between model quality, training computation, and inference cost\. KITE provides a common setting for studying these choices without fixing the architecture to the particular connections used in SST\.

## 6Related work

Net2Net studies function\-preserving model growth\[[6](https://arxiv.org/html/2609.27294#bib.bib3)\]; SPARKLING examines signal preservation and symmetry breaking in width\-progressive learning\[[28](https://arxiv.org/html/2609.27294#bib.bib4)\]\. Dense\-to\-MoE upcycling reuses dense checkpoints to initialize sparse expert models\[[16](https://arxiv.org/html/2609.27294#bib.bib14),[10](https://arxiv.org/html/2609.27294#bib.bib6)\]\. SST instead adds the Decoder to an already sparse source partway through training\. Its aligned initialization changes the prediction graph and is followed by joint adaptation\.

Multi\-query attention\[[20](https://arxiv.org/html/2609.27294#bib.bib15)\]and grouped\-query attention\[[1](https://arxiv.org/html/2609.27294#bib.bib16)\]reduce KV storage by sharing KV heads across query heads\. Cross\-layer attention extends sharing across adjacent layers\[[4](https://arxiv.org/html/2609.27294#bib.bib17)\]\. These approaches reduce the stored attention state; KITE targets capacity expansion without extending the KV\-producing path\.

YOCO separates cache\-producing and cache\-consuming computation and exploits this separation for prefill early exit\[[22](https://arxiv.org/html/2609.27294#bib.bib7)\]\. SST applies this execution pattern to staged model growth\. YOCO’s cross\-decoder shares a global KV representation; SST reads aligned layer\-wise KV and is constructed through source\-to\-target expansion\.

Mixture\-of\-Recursions explores recursive KV sharing, reusing the first recursion’s KV across subsequent recursions\[[3](https://arxiv.org/html/2609.27294#bib.bib11)\]\. From an architectural perspective, such fixed\-KV looping can be viewed as a parameter\-tied special case of the family underlying KITE\. SST instead uses independent Prefiller and Decoder parameters, expanding capacity during training without extending the KV\-producing path\.

## 7Conclusion

KITE expands a model during training by adding capacity that reuses KV, leaving bulk\-prefill computation source\-sized\. SST provides a simple instance of this route\. With both training stages counted, it reaches lower training loss than the 47B and 63B classic baselines at comparable theoretical training compute\. It also scores higher on several downstream tasks\. Relative to the 47B baseline, SST trades lower bulk\-prefill computation for higher decode computation; its inference\-cost proxy is lower in prefill\-heavy workloads\. The larger 63B reference has a higher inference proxy and slightly higher training loss\. These observations support further investigation of the combined scaling route, without establishing a scaling law, a causal benefit from each ingredient, or measured serving speedup\.

## Appendix AOpenRouter workload statistics and inference proxy

#### Data and coverage\.

Table[1](https://arxiv.org/html/2609.27294#S1.T1)summarizes six selected models using public OpenRouter daily activity retrieved on September 21, 2026\. We aggregate available standard and Batch activity within August 22–September 20 UTC, without extrapolating missing days\. Both Astra variants have 17 observed days \(September 4–20\), and Gemini 3\.8 Flash has 19 \(September 2–20\)\. The other models have 30 standard\-activity days; Claude Fable 5 has only 10 Batch days \(August 22–31\)\. This selected traffic is not separately labeled as agentic usage\.

#### Token and charge accounting\.

LetPPbe total input tokens,CCcache\-read tokens,U=P−CU=P\-Cuncached input tokens, andDDoutput tokens, including reasoning\. The two token ratios areP/DP/DandU/DU/D\. For each variantvv, letava\_\{v\}andovo\_\{v\}be the September 21 base input and output prices per million tokens\. We estimate

A=10−6​∑vUv​av,O=10−6​∑vDv​ov,suncached=AA\+O\.A=10^\{\-6\}\\sum\_\{v\}U\_\{v\}a\_\{v\},\\qquad O=10^\{\-6\}\\sum\_\{v\}D\_\{v\}o\_\{v\},\\qquad s\_\{\\mathrm\{uncached\}\}=\\frac\{A\}\{A\+O\}\.\(6\)Standard and Batch charges are calculated separately before summation\. Cache\-read charges are excluded from both numerator and denominator\. These are base\-price estimates, not invoices; cache\-write premiums, context\-length and provider\-specific pricing, historical price changes, and non\-token fees are not modeled\.

Table 4:Components underlying Table[1](https://arxiv.org/html/2609.27294#S1.T1)\. Token counts are in billions; estimated charges are in millions of USD\.UUis uncached input,CCcached input, andDDoutput\. Ratios use unrounded totals\.Standard input/output rates \(USD per million tokens\) are 10/50 for Astra, Astra Pro and Fable 5; 2/10 for Sonnet 5; and 0\.75/3\.75 for both Gemini models\. Batch rates are half the corresponding standard rates\.

#### Sources\.

#### Relation to the inference proxy\.

These statistics motivate input\-heavy inference but do not measure GPU costs\. We apply the same prefill:decode cost weights to all models\. For prefill weightww, the proxy isw​p\+\(1−w\)​dwp\+\(1\-w\)d, whereppandddare bulk\-prefill and decode active\-body parameter counts normalized to 47B classic\. Thusw=0\.75w=0\.75means a 75:25 P:D cost mix\. Figure[1](https://arxiv.org/html/2609.27294#S0.F1)uses this mix, giving 0\.933 for SST and 1\.365 for 63B classic, compared with 1 for 47B classic: SST is 6\.7% below 47B and 31\.6% below 63B\. Figure[5](https://arxiv.org/html/2609.27294#S4.F5)varies the same cost weights\. These are analytical proxy comparisons, not measured serving speedups\.

## Appendix BParameter accounting

Parameter counts follow the configured tensor shapes, checked against the exported weight inventory\. SST has two sets of Transformer blocks and one shared embedding, output norm, and output head\. Only the Prefiller contains K/V projections and K normalization in the effective architecture\. The entry and final\-state RMS operations have no trainable parameters\.

Table 5:Detailed geometry\. Source/SST dimensions apply to each tower; both have two initial dense FFN layers\. RoPE denotes rotary position embeddings\[[21](https://arxiv.org/html/2609.27294#bib.bib19)\]\.Total parameters count the effective architecture\. Active body parameters exclude embedding/head and count Top\-8 routed experts per MoE layer, shared experts, routers, and the remaining body weights needed for one token\. These are per\-token parameter counts, not measured execution costs\.

Table 6:Architectural parameter counts in billions, rounded to three decimal places\.The exported vocabulary has 128,815 rows, without training padding\. Excluding embedding and head from the active\-body counts avoids treating the full embedding matrix as a per\-token multiplication\. The output head still contributes prompt\-boundary and decode work and must be included in measured serving cost\.

## Appendix CTraining budgets and loss reporting

Table 7:Training budgets at the final checkpoints\. Theoretical training FLOPs are normalized to the 47B classic reference\.SST’s total includes 167\.98B source\-stage tokens and 222\.55B continuation tokens; both stages are included in its training compute\.

The loss curves end at the evaluated final checkpoints\. Training\-loss curves use an exponential moving average with span 200 \(adjust=False\), computed from the recorded observations without interpolation\. The EMA is restarted after SST conversion and computed before plot subsampling, with the final observation retained\.

The evaluated 47B classic checkpoint definesBref≈5\.16836×1021B\_\{\\mathrm\{ref\}\}\\approx 5\.16836\\times 10^\{21\}FLOPs\. Forward costs are approximately 3\.887488B FLOPs/token for 47B classic, 5\.148513B for 63B classic, 3\.087252B for the source, and 5\.410678B for SST\. SST retains both towers’ attention and FFN/MoE computation but only one embedding, output head, and set of KV projections\.

All forward coefficients use sequence length 4096, sliding\-window size 512, and the padded training vocabulary of 128,896, with multiply\-add counted as two FLOPs\. They include attention projections and matrix products, embedding, output head, FFN/MoE, routers, and the estimator’s activation terms\. Normalization, RoPE, residual and bridge operations, rematerialization, optimizer, loss/auxiliary loss, conversion, communication, and checkpoint I/O are excluded\. The same theoretical accounting convention is used for all models\.

## Appendix DTraining recipe and continuation scheduler

Continued training uses native BF16 parameters with Muon and Adam groups\. Muon uses six Newton–Schulz steps, with packed attention split by head and GLU split by projection type; Adam usesϵ=10−15\\epsilon=10^\{\-15\}\. The source peak base learning rate is approximately1\.03302×10−31\.03302\\times 10^\{\-3\}\. The continuation scheduler uses a base weight decay of 0\.15 with the inherited parameter\-group multipliers\. The source/SST/63B token batch is 12,582,912; the 47B token batch is 14,680,064\.

SST continuation uses 128 NVIDIA H800 GPUs with expert parallelism of size 8, a sequence length of 4,096, and a global batch size of 3,072 sequences\. It runs for 17,687 updates after 13,350 source\-training updates, giving 31,037 updates in total\.

The parameter\-free RMS operations at the Decoder input and final state accumulate in FP32 and return the input dtype; the original learned output norm is applied after the final\-state RMS operation\.

Letuube the number of continuation tokens consumed,U=222,553,964,544U=222\{,\}553\{,\}964\{,\}544the configured continuation horizon, andq0q\_\{0\}the source’s normalized cosine phase at conversion\. SST uses

q⁡\(u\)\\displaystyle q\(u\)=q0\+\(1−q0\)​min⁡\(u/U,1\),\\displaystyle=q\_\{0\}\+\(1\-q\_\{0\}\)\\min\(u/U,1\),\(7\)η⁡\(u\)\\displaystyle\\eta\(u\)=ηmin\+\(ηanchor−ηmin\)​1\+cos⁡\(π​q​\(u\)\)1\+cos⁡\(π​q0\)\.\\displaystyle=\\eta\_\{\\min\}\+\(\\eta\_\{\\mathrm\{anchor\}\}\-\\eta\_\{\\min\}\)\\frac\{1\+\\cos\(\\pi q\(u\)\)\}\{1\+\\cos\(\\pi q\_\{0\}\)\}\.\(8\)The source schedule warms up over 12\.582912B tokens and decays until 335\.9605248B tokens\. Conversion occurs at 167\.9818752B tokens\. Thusq0q\_\{0\}is measured within the post\-warmup decay interval, not by dividing the conversion position by the whole horizon\. The learning\-rate constants are

ηanchor\\displaystyle\\eta\_\{\\mathrm\{anchor\}\}≈5\.96548×10−4,\\displaystyle\\approx 5\.96548\\times 10^\{\-4\},ηmin\\displaystyle\\eta\_\{\\min\}≈1\.03302×10−4\.\\displaystyle\\approx 1\.03302\\times 10^\{\-4\}\.Parameter\-group multipliers remain applicable to the common base learning rate\.

## References

- \[1\]\(2023\)GQA: Training Generalized Multi\-Query Transformer Models from Multi\-Head Checkpoints\.arXiv preprint arXiv:2305\.13245\.External Links:[Link](https://arxiv.org/abs/2305.13245)Cited by:[§6](https://arxiv.org/html/2609.27294#S6.p2.1)\.
- \[2\]J\. Austin, A\. Odena, M\. Nye, M\. Bosma, H\. Michalewski, D\. Dohan, E\. Jiang, C\. Cai, M\. Terry, Q\. Le,et al\.\(2021\)Program Synthesis with Large Language Models\.arXiv preprint arXiv:2108\.07732\.External Links:[Link](https://arxiv.org/abs/2108.07732)Cited by:[§4\.2](https://arxiv.org/html/2609.27294#S4.SS2.p1.1)\.
- \[3\]S\. Bae, Y\. Kim, R\. Bayat, S\. Kim, J\. Ha, T\. Schuster, A\. Fisch, H\. Harutyunyan, Z\. Ji, A\. C\. Courville, and S\. Yun\(2025\)Mixture\-of\-Recursions: learning dynamic recursive depths for adaptive token\-level computation\.InAdvances in Neural Information Processing Systems,Vol\.38\.External Links:[Link](https://arxiv.org/abs/2507.10524)Cited by:[§6](https://arxiv.org/html/2609.27294#S6.p4.1)\.
- \[4\]W\. Brandon, M\. Mishra, A\. Nrusimha, R\. Panda, and J\. Ragan\-Kelly\(2024\)Reducing Transformer Key\-Value Cache Size with Cross\-Layer Attention\.arXiv preprint arXiv:2405\.12981\.External Links:[Link](https://arxiv.org/abs/2405.12981)Cited by:[§6](https://arxiv.org/html/2609.27294#S6.p2.1)\.
- \[5\]M\. Chen, J\. Tworek, H\. Jun, Q\. Yuan, H\. P\. d\. O\. Pinto, J\. Kaplan, H\. Edwards, Y\. Burda, N\. Joseph, G\. Brockman,et al\.\(2021\)Evaluating Large Language Models Trained on Code\.arXiv preprint arXiv:2107\.03374\.External Links:[Link](https://arxiv.org/abs/2107.03374)Cited by:[§4\.2](https://arxiv.org/html/2609.27294#S4.SS2.p1.1)\.
- \[6\]T\. Chen, I\. Goodfellow, and J\. Shlens\(2016\)Net2Net: accelerating learning via knowledge transfer\.InInternational Conference on Learning Representations,External Links:[Link](https://arxiv.org/abs/1511.05641)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p3.1),[§6](https://arxiv.org/html/2609.27294#S6.p1.1)\.
- \[7\]K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano,et al\.\(2021\)Training Verifiers to Solve Math Word Problems\.arXiv preprint arXiv:2110\.14168\.External Links:[Link](https://arxiv.org/abs/2110.14168)Cited by:[§4\.2](https://arxiv.org/html/2609.27294#S4.SS2.p1.1)\.
- \[8\]DeepSeek\(2026\)DeepSeek\-V4\.1\-Flash: pushing the limits of kv cache compression\.Note:DeepSeek\-V4\.1\-Flash,[https://huggingface\.co/deepseek\-ai/DeepSeek\-V4\.1\-Flash](https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash)Technical report and model card\. Accessed September 22, 2026\.Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p3.1)\.
- \[9\]W\. Fedus, B\. Zoph, and N\. Shazeer\(2022\)Switch Transformers: scaling to trillion parameter models with simple and efficient sparsity\.Journal of Machine Learning Research23\(120\),pp\. 1–39\.External Links:[Link](https://arxiv.org/abs/2101.03961)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p2.1)\.
- \[10\]E\. He, A\. Khattar, R\. Prenger, V\. Korthikanti, Z\. Yan, T\. Liu, S\. Fan, A\. Aithal, M\. Shoeybi, and B\. Catanzaro\(2024\)Upcycling large language models into mixture of experts\.arXiv preprint arXiv:2410\.07524\.External Links:[Link](https://arxiv.org/abs/2410.07524)Cited by:[§6](https://arxiv.org/html/2609.27294#S6.p1.1)\.
- \[11\]D\. Hendrycks, C\. Burns, S\. Basart, A\. Zou, M\. Mazeika, D\. Song, and J\. Steinhardt\(2020\)Measuring Massive Multitask Language Understanding\.arXiv preprint arXiv:2009\.03300\.External Links:[Link](https://arxiv.org/abs/2009.03300)Cited by:[§4\.2](https://arxiv.org/html/2609.27294#S4.SS2.p1.1)\.
- \[12\]D\. Hendrycks, C\. Burns, S\. Kadavath, A\. Arora, S\. Basart, E\. Tang, D\. Song, and J\. Steinhardt\(2021\)Measuring Mathematical Problem Solving With the MATH Dataset\.arXiv preprint arXiv:2103\.03874\.External Links:[Link](https://arxiv.org/abs/2103.03874)Cited by:[§4\.2](https://arxiv.org/html/2609.27294#S4.SS2.p1.1)\.
- \[13\]J\. Hoffmann, S\. Borgeaud, A\. Mensch, E\. Buchatskaya, T\. Cai, E\. Rutherford, D\. de Las Casas, L\. A\. Hendricks, J\. Welbl, A\. Clark,et al\.\(2022\)Training compute\-optimal large language models\.arXiv preprint arXiv:2203\.15556\.External Links:[Link](https://arxiv.org/abs/2203.15556)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p1.1)\.
- \[14\]J\. Kaplan, S\. McCandlish, T\. Henighan, T\. B\. Brown, B\. Chess, R\. Child, S\. Gray, A\. Radford, J\. Wu, and D\. Amodei\(2020\)Scaling laws for neural language models\.arXiv preprint arXiv:2001\.08361\.External Links:[Link](https://arxiv.org/abs/2001.08361)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p1.1)\.
- \[15\]D\. P\. Kingma and J\. Ba\(2014\)Adam: A Method for Stochastic Optimization\.arXiv preprint arXiv:1412\.6980\.External Links:[Link](https://arxiv.org/abs/1412.6980)Cited by:[§3\.2](https://arxiv.org/html/2609.27294#S3.SS2.p2.1)\.
- \[16\]A\. Komatsuzaki, J\. Puigcerver, J\. Lee\-Thorp, C\. R\. Ruiz, B\. Mustafa, J\. Ainslie, Y\. Tay, M\. Dehghani, and N\. Houlsby\(2022\)Sparse Upcycling: Training Mixture\-of\-Experts from Dense Checkpoints\.arXiv preprint arXiv:2212\.05055\.External Links:[Link](https://arxiv.org/abs/2212.05055)Cited by:[§6](https://arxiv.org/html/2609.27294#S6.p1.1)\.
- \[17\]J\. Liu, J\. Su, X\. Yao, Z\. Jiang, G\. Lai, Y\. Du, Y\. Qin, W\. Xu, E\. Lu, J\. Yan,et al\.\(2025\)Muon is Scalable for LLM Training\.arXiv preprint arXiv:2502\.16982\.External Links:[Link](https://arxiv.org/abs/2502.16982)Cited by:[§3\.2](https://arxiv.org/html/2609.27294#S3.SS2.p2.1)\.
- \[18\]T\. Mihaylov, P\. Clark, T\. Khot, and A\. Sabharwal\(2018\)Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering\.arXiv preprint arXiv:1809\.02789\.External Links:[Link](https://arxiv.org/abs/1809.02789)Cited by:[§4\.2](https://arxiv.org/html/2609.27294#S4.SS2.p1.1)\.
- \[19\]N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. Le, G\. Hinton, and J\. Dean\(2017\)Outrageously Large Neural Networks: The Sparsely\-Gated Mixture\-of\-Experts Layer\.arXiv preprint arXiv:1701\.06538\.External Links:[Link](https://arxiv.org/abs/1701.06538)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p2.1)\.
- \[20\]N\. Shazeer\(2019\)Fast Transformer Decoding: One Write\-Head is All You Need\.arXiv preprint arXiv:1911\.02150\.External Links:[Link](https://arxiv.org/abs/1911.02150)Cited by:[§6](https://arxiv.org/html/2609.27294#S6.p2.1)\.
- \[21\]J\. Su, Y\. Lu, S\. Pan, A\. Murtadha, B\. Wen, and Y\. Liu\(2021\)RoFormer: Enhanced Transformer with Rotary Position Embedding\.arXiv preprint arXiv:2104\.09864\.External Links:[Link](https://arxiv.org/abs/2104.09864)Cited by:[Table 5](https://arxiv.org/html/2609.27294#A2.T5)\.
- \[22\]Y\. Sun, L\. Dong, Y\. Zhu, S\. Huang, W\. Wang, S\. Ma, Q\. Zhang, J\. Wang, and F\. Wei\(2024\)You only cache once: decoder\-decoder architectures for language models\.arXiv preprint arXiv:2405\.05254\.External Links:[Link](https://arxiv.org/abs/2405.05254)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p3.1),[§6](https://arxiv.org/html/2609.27294#S6.p3.1)\.
- \[23\]M\. Suzgun, N\. Scales, N\. Schärli, S\. Gehrmann, Y\. Tay, H\. W\. Chung, A\. Chowdhery, Q\. V\. Le, E\. H\. Chi, D\. Zhou,et al\.\(2022\)Challenging BIG\-Bench Tasks and Whether Chain\-of\-Thought Can Solve Them\.arXiv preprint arXiv:2210\.09261\.External Links:[Link](https://arxiv.org/abs/2210.09261)Cited by:[§4\.2](https://arxiv.org/html/2609.27294#S4.SS2.p1.1)\.
- \[24\]A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. Polosukhin\(2017\)Attention Is All You Need\.arXiv preprint arXiv:1706\.03762\.External Links:[Link](https://arxiv.org/abs/1706.03762)Cited by:[§2\.2](https://arxiv.org/html/2609.27294#S2.SS2.p1.1)\.
- \[25\]WeChat AI\(2026\)Building effective sparse MoE models with moderate resources\.Note:WeLM Blog,[https://welm\.weixin\.qq\.com/en/posts/building\-effective\-sparse\-moe\-models\-with\-moderate\-resources/](https://welm.weixin.qq.com/en/posts/building-effective-sparse-moe-models-with-moderate-resources/)Published January 21, 2026\. Source of the KV\-Mirror architecture\.Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p3.1)\.
- \[26\]J\. Yang, C\. E\. Jimenez, A\. Wettig, K\. Lieret, S\. Yao, K\. Narasimhan, and O\. Press\(2024\)SWE\-agent: Agent\-Computer Interfaces Enable Automated Software Engineering\.arXiv preprint arXiv:2405\.15793\.External Links:[Link](https://arxiv.org/abs/2405.15793)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p5.1)\.
- \[27\]S\. Yao, J\. Zhao, D\. Yu, N\. Du, I\. Shafran, K\. Narasimhan, and Y\. Cao\(2022\)ReAct: synergizing reasoning and acting in language models\.arXiv preprint arXiv:2210\.03629\.External Links:[Link](https://arxiv.org/abs/2210.03629)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p5.1)\.
- \[28\]Q\. Yu, X\. Ma, Z\. Zhuo, M\. Wang, D\. Liu, S\. Zhan, Y\. Ma, L\. Xiang, X\. Bin, and D\. He\(2026\)SPARKLING: balancing signal preservation and symmetry breaking for width\-progressive learning\.arXiv preprint arXiv:2602\.02472\.External Links:[Link](https://arxiv.org/abs/2602.02472)Cited by:[§1](https://arxiv.org/html/2609.27294#S1.p3.1),[§6](https://arxiv.org/html/2609.27294#S6.p1.1)\.