PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashNext-RL: A System Optimization Study for Long-Horizon Agentic RL of Qwen3.8-Flash-Next

Zhewei Yao, Yuxiong He

arXiv:2610.28998v1cs.LGcs.CL

TL;DR

Long-horizon agentic RL of Qwen3.8-Flash-Next is challenging because its memory requirements and model-specific operators exceed what evaluated open-source implementations support at the target context. FlashNext-RL combines model-specific BF16 training optimizations with vLLM-based SWE rollout generation. The resulting system supports a target batch of 1,024 trainable rollouts using one training node and three inference nodes.

  • Problem

    Evaluated open-source implementations did not meet the target of efficient Qwen3.8-Flash-Next training at 131K-token context on one B200 node, amid substantial activation and host-memory demands.

    Source: §1 · p. 2

  • Method

    FlashNext-RL combines a single-node BF16 trainer with vLLM-based SWE rollout generation and model-specific training optimizations for Qwen3.8-Flash-Next.

    Source: §1 · p. 2

  • Results

    The system supports a target batch of 1,024 trainable rollouts with one training node and three inference nodes.

    Source: §6 · p. 21

  • Takeaways & Limitations

    The study’s objective is to identify model-specific RL training optimizations that make long-horizon Qwen3.8-Flash-Next RL practical, rather than to optimize accuracy or recipe efficiency.

    Source: §1 · p. 2

    The work does not aim to establish the best model accuracy or the most efficient RL recipe.

    Source: §1 · p. 2

Abstract

from arXiv · show

Long-horizon agentic reinforcement learning (RL) requires both efficient rollout generation and training on long sequences. However, the memory requirements and model-specific operators of Qwen3.8-Flash-Next make this difficult with existing open-source implementations. To resolve these limitations, we introduce FlashNext-RL, a model-specific RL training optimization framework based on prime-rl, with vLLM for software engineering (SWE) rollout generation. Our main focus is BF16 training of the full 48-layer model on a single 8xB200 node, with the first two layers, the n-gram table, the sparse-attention indexer and the multi-token prediction head frozen. We optimize host offload, gradient accumulation and the kernels for mixture-of-experts, sparse attention and the gated residual stream. In the full-model development study, non-padding training throughput increases from approximately 395 tokens/s/GPU in the initial 32K-token configuration to 4,438 tokens/s/GPU in the optimized 131K-token configuration, with hardware matmul throughput (aka without considering the cost of attention components) increasing from approximately 20 to 213 TFLOPs/GPU. A controlled cumulative study on an 8-layer model further quantifies the contributions of individual optimizations. Together with three inference nodes, the system supports a live SWE RL run with a target batch of 1,024 trainable rollouts and a 131,072-token context limit. Over steps 5-15, an iteration takes 22.4 minutes on average, with bounded entropy and trainer-sampler mismatch KL. These results highlight the importance of tailoring RL training to the model's architecture, memory layout and kernel requirements, rather than relying only on general-purpose RL infrastructure.

1 Introduction

FlashNext-RL targets long-horizon SWE reinforcement learning of Qwen3.8-Flash-Next, whose memory demands and model-specific operators challenge open-source training systems. It combines model-specific training optimizations with vLLM rollout generation, while focusing on practicality rather than model accuracy or the most efficient RL recipe.

1 Introduction

  • Without context parallelism, retaining one residual-stream input per layer for activation checkpointing requires about 129 GB (120 GiB) per GPU for 48 × 131,072 × 10,240 BF16 values, before other memory.With CPU offload, parameter, gradient and AdamW state is budgeted at about 1.9 TiB before the n-gram table and remaining RL system state, exceeding node host-memory capacity once included. Few efficient implementations were found for the sparse grouped-query attention backward pass, intermediate-size-640 MoE kernels at expert parallelism 8, or fused gated-residual and per-layer-embedding kernels.Source: §1 · p. 2
  • FlashNext-RL combines a single-node BF16 trainer with vLLM rollout generation for Qwen3.8-Flash-Next, using SWE as its agentic-RL example.The full 48-layer model is trained with the first two layers, n-gram table, QSA indexer, and MTP head frozen.Source: §1 · p. 2
  • The training configuration reaches 4,438 non-padding tokens/s/GPU and 213 hardware matmul TFLOPs/GPU in the optimized full-model development study.Isolated comparisons and an 8-layer cumulative study quantify individual optimization contributions.Source: §1 · p. 2
  • The system combines one training node with three inference nodes, using memory, weight-transfer, and checkpoint optimizations for the RL loop.The introduction identifies an end-to-end SWE RL run reporting iteration time, throughput, sequence lengths, and short-run numerical behavior.Source: §1 · p. 2
  • The work’s objective is to identify model-specific training optimizations that make long-horizon Qwen3.8-Flash-Next RL practical, not to maximize accuracy or recipe efficiency.The model uses hybrid attention, gated residuals, n-gram embeddings, fine-grained MoE, and an MTP head.Source: §1 · pp. 1–2

2 Inference Optimization through vLLM Configuration Tuning

The inference section configures vLLM to serve long SWE rollouts and selects a tensor/expert-parallel topology against throughput and KV-cache capacity. Its live behavior reports output throughput, session activity, cache use, and MTP acceptance.

2.1 Topology

  • Expert parallelism avoids padding 640-wide experts to 256 columns per rank, reducing routed-expert memory from 90.0 GiB to 56.25 GiB per rank.TP padding reduces KV-cache capacity and increases selected-expert weight traffic during decoding.Source: §2.1 · p. 3
  • TP 4 with expert parallelism is selected for target concurrency, reaching 1,091 output tokens/s/GPU at 26 sessions per GPU in replay.TP 2 is fastest at light load but its prefix cache stops hitting beyond 10.5 sessions per GPU for 64K-token SWE contexts; TP 8 has the lowest throughput at every load.Source: §2.1 · p. 3, Table 1 · p. 3

2.2 Final Configuration

  • The final serving setup uses two TP 4 plus EP engines per inference node, BF16 weights and KV cache, prefix caching, and MTP3 speculative decoding.The MTP head accepts 1.82–1.88 draft tokens per round on replayed rollouts, but is not updated during training.Source: §2.2 · p. 3
  • A shared node-local tmpfs copy of the frozen n-gram table replaces two engine copies, with the second engine adding only 2.2 GB.The table must reside on tmpfs because cudaHostRegister rejects a writable shared mapping of a regular file.Source: §2.2 · p. 4
  • Router replay returns each token’s routed experts from vLLM to the trainer, addressing implementation differences that explain about 96% of offline trainer–inference mismatch.The payload is [T, 48, 10] uint16, or 960 bytes per token; live mismatch KL is reported separately.Source: §2.2 · p. 4

2.3 Live Behavior

  • Live-run telemetry reports 17.1K output tokens/s across six engines, while about 77 of 112 sessions per engine are generating at a time.Requests wait only 0.4–0.8 on average for an engine slot.Source: §2.3 · p. 4
  • KV-cache use averages 50% and peaks at 72%, with prefix-cache hits on 85% of prompt tokens.The admission cap engages only at cache-use peaks.Source: §2.3 · p. 4
  • The live run’s mean MTP acceptance length is 2.55 tokens per target-model forward pass.Source: §2.3 · p. 4

3 Training Optimization

Training optimization focuses on BF16 full-model training at a 131K-token bin on one 8×B200 node, then uses an 8-layer model to study cumulative optimization effects.

3 Training Optimization

  • The target is BF16 training of the full 48-layer model at a 131K-token bin on one 8×B200 node, with cumulative optimization effects evaluated on an 8-layer model.The full-model configuration freezes components listed earlier, while the 8-layer ladder does not freeze its first two layers.Source: §3 · p. 4

3.1 Setup

The setup defines replay-based throughput and FLOP accounting for long SWE sequences, then lays out memory constraints and the training parallel/offload design. Its budget motivates a smaller optimizer footprint alongside CPU offload and sequence parallelism.

3.1 Setup

  • Throughput replays recorded SWE steps without generation or weight broadcast, packs batches into 131,072-token bins, and reports non-padding tokens per second per GPU.The cumulative ladder disables PSS sampling because page-table walks slow this offload-heavy workload; its memory ladder step times are only indicative.Source: §3.1 · pp. 4–5
  • For the full model with two layers frozen, the accounted hardware matmul cost is 48.4 GFLOP per token, excluding attention, indexer scores, and GDN recurrence.TFLOP accounting includes checkpoint recompute, attention-backward score recompute, and output-head logit recompute; the output head uses fp32/TF32 and GDN retains fp32 state.Source: §3.1 · p. 5
  • The host-memory plan combines CPU-offloaded parameters and optimizer state with Muon for matrix parameters, because AdamW’s two moments do not fit alongside the n-gram table and system state.FSDP2 gathers units on GPU as needed, while expert gradients accumulate on GPU before offload; StreamingMuon streams parameters from host memory for GPU orthogonalization.Source: §3.1 · p. 5, Table 2 · p. 5
  • The layout assigns 64 of each layer’s 512 experts to each GPU and splits 131K sequences across four GPUs with Ulysses context parallelism.The 8-layer model keeps the full output head, which accounts for 19% of its FLOPs versus 4% at 48 layers, so their TFLOPs/GPU are not comparable.Source: §3.1 · p. 5
  • Full activation checkpointing reuses forward-pass routing decisions during recomputation, and the first two layers are frozen in the full-model configuration.Source: §3.1 · p. 5

3.2 Isolated Optimizations

The isolated studies address host-memory limits and host-copy stalls while testing context-parallel layouts for 131K-token training. Pinned offload, GPU-resident expert-gradient accumulation, and CP 4 improve throughput under the reported settings.

3.2.1 Host Offload: Pageable vs. PinnedArena

  • The Muon step drops from 66–75 s to 12.4 s with PinnedArena, which uses fixed-size registered host-memory regions.The default pinned allocator rounded blocks up by 28% and was killed at 1,887 GiB; pageable offload fit at 1,555 GiB but copied more slowly.Source: §3.2.1 · p. 6, Table 3 · p. 6

3.2.2 Gradients: Blocking Host Copies vs. GPU-Resident Expert Gradients

  • A 2.4–2.7× throughput improvement follows GPU-resident accumulation of EP-local expert gradients, with dense accumulation adding 4% on the final configuration.Expert gradients transfer to the host once per step; the dense-accumulator comparison is 3,233/3,277 →3,383/3,372 REAL tokens/s/GPU.Source: §3.2.2 · p. 6, Table 4 · p. 6

3.2.3 Context Parallelism: Ulysses CP 4 vs. CP 8

  • 1,368/1,571 vs. 1,132/1,182 REAL tokens/s/GPU favors CP 4 over CP 8, and CP 4 fits despite its higher peak allocation.CP 8 peaks at 124.9 vs. 149.3 GiB allocated for CP 4; grouping checkpointed layers would save memory but cost 11–22% more backward time.Source: §3.2.3 · p. 7

3.2.4 MoE: DeepEP vs. Mixture-of-Kittens

The isolated MoE, expert-placement, and QSA-indexer studies target costly model-specific operations. Fused MoE execution improves layer speed and memory use, while placement and indexer changes provide additional measured gains.

3.2.4 MoE: DeepEP vs. Mixture-of-Kittens

  • 1.56–2.24× faster per layer, Mixture-of-Kittens fuses MoE dispatch, expert GEMMs, and combine in persistent forward and backward kernels.At 131K tokens per rank, its peak is 15.9 vs. 71.1 GiB, though symmetric-memory workspace may be outside PyTorch’s allocator accounting.Source: §3.2.4 · p. 7, Table 5 · p. 8
  • 5.3% higher full-model throughput (3,405 →3,586 REAL tokens/s/GPU) follows native intermediate-size support and increased communication SMs.The native size reduces layer time by 4.8%; the subsequent SM change reduces the remaining time by 18.7%.Source: §3.2.4 · p. 7

3.2.5 Static Expert Placement

  • The max/mean rank load falls from 1.233 to 1.059 on the full model with static expert placement, but its final-stack throughput effect is within 1%.Placement gave +3.4% on 8 layers; the hottest expert carries at most 0.2 of a rank’s mean load.Source: §3.2.5 · p. 8

3.2.6 QSA Indexer: PyTorch vs. cuDNN Fused Indexer + Top-k

  • 5.5× faster at 131K, the fused cuDNN indexer plus top-k replaces PyTorch after zero-padding QSA’s four indexer heads to 32.The frozen indexer’s selection is computed once in the forward pass and reused during recompute.Source: §3.2.6 · p. 8, Table 6 · p. 8

3.2.7 QSA Attention: Dense vs. cuDNN DSA vs. Triton vs. CuTe

QSA requires token-level sparse attention because dense attention computes a different function beyond 2,048 tokens. The study compares sparse-attention implementations and reports full-model gains from TMA gathers and a CuTe kernel, alongside a fused gated-residual approach.

3.2.7 QSA Attention: Dense vs. cuDNN DSA vs. Triton vs. CuTe

  • Sparse attention is required: dense QSA computes a different function beyond 2,048 tokens, while block-sparse masks touch 90–94% of causal key tiles on real data.The cuDNN DSA implementation leaves 52 of 64 forward head slots unused and its backward pass reaches about 95 TFLOP/s; GQA-native Triton is 1.58× faster than DSA at 16K and 1.34× at 64K.Source: §3.2.7 · pp. 8–9
  • The CuTe kernel is 1.63× faster than Triton TMA at 64K, and TMA gathers plus CuTe raise full-model throughput from 3,412 to 3,642 REAL tokens/s/GPU.The kernel takes 46.7 ms at 64K, near the gather and MMA floors of about 41 and 40 ms.Source: §3.2.7 · pp. 8–9, Table 7 · p. 9

3.2.8 Gated Residual: Eager vs. Compiled vs. Fused

  • 1.95× speedup over torch.compile at 131K, fused gated-residual read and inject keeps 8× less memory per pair and cuts real GDN-layer forward-plus-backward time by 31%.The kernels recompute cheap intermediates during backward, lowering transient memory per layer by about 10 GiB.Source: §3.2.8 · p. 9, Table 8 · p. 10

3.2.9 Other Fused Kernels

The remaining optimizations fuse kernels, move n-gram lookup to a GPU-accessible host mapping, address compilation stalls, and freeze the first two layers. Their reported gains vary by operation and model scale.

3.2.9 Other Fused Kernels

  • +6.4% on 8 layers and +2.7% on the full model come from fused attention q/k norm and partial RoPE; decoder glue adds +0.7% on the full model.The q/k operator takes 2.28 ms versus 13.5 ms eager and 5.2 ms compiled at 32K tokens and 24 heads.Source: §3.2.9 · p. 10

3.2.10 N-gram Lookup: CPU Gather vs. UVA Gather

  • +8.1% full-model throughput (3,402 →3,688 REAL tokens/s/GPU) follows GPU-side UVA gathering of the frozen host-resident n-gram table.CPU index_select left the GPU idle for about 540 ms of every micro step, or 6.3%.Source: §3.2.10 · p. 10, Table 9 · p. 10

3.2.11 Compilation and Stalls

  • 8-layer throughput rises from 5,059 to 8,243 REAL tokens/s/GPU after removing token count from FLA autotuning keys and applying other concurrent fixes.One rank had been idle 45.5% of the time during retuning; pre-warming cuts step 1 from 171.9 to 23.8 s without changing steady state.Source: §3.2.11 · pp. 10–11

3.2.12 Freezing the First Two Layers

  • Freezing the first two layers is worth about 10% full-model throughput and 37 GiB of host memory, while reducing the 8-layer activation peak from 107.6 to 78.6 GiB.Freezing removes those layers’ gradients, optimizer state, and backward pass, including the PLE layer’s large transient.Source: §3.2.12 · p. 11

3.3 Putting It Together

The full-model development study reaches 4,438 REAL tokens/s/GPU at a 131K-token bin, while an 8-layer cumulative study measures gains from individual optimizations. The 8-layer results show that offload improvements and model-specific kernels contribute at different stages of the optimization stack.

8-layer ladder

  • 2,391 to 17,678 REAL tokens/s/GPU (7.4×) and 20 to 146 model TFLOPs/GPU (195 including recompute) mark the 8-layer cumulative ladder’s endpoints.The ladder adds optimizations one at a time to the first eight layers on real rollouts.Source: §3.3 · p. 11, Figure 1 · p. 12, Table 10 · p. 12
  • Keeping parameters resident across micro steps adds 64% throughput, and keeping expert gradients on GPU adds 2.54×, making the offload path the largest source of ladder gains.Before these changes, GPUs mostly wait for host copies.Source: §3.3 · p. 11
  • MoK adds 24.6%, the fused gated residual 14%, GQA-native QSA 13.6%, UVA n-gram gather 5.6%, CuTe dK/dV 3.4%, and fused glue 3.2%.These kernel gains are reported as sequential additions to the ladder.Source: §3.3 · p. 11
  • torch.compile costs 20% on L0 but gains 2% on the final stack, while expert placement, pre-warming, and expert prefetch have small or negligible effects at eight layers.Removing MoK, the fused gated residual, or sparse-attention kernels individually from the final stack costs 15–22%.Source: §3.3 · p. 11

Full model

  • 4,438 REAL tokens/s/GPU and 213 hardware matmul TFLOPs/GPU are reached by the optimized 48-layer configuration at a 131K-token bin.The initial configuration records approximately 395 REAL tokens/s/GPU at a 32K-token bin; settings vary across rows, so this is not a controlled speedup comparison.Source: §3.3 · p. 11

8-layer ladder

  • The ladder includes correctness and packing fixes that cannot be turned off, including power-of-two indexer buckets and reuse of forward routes during recompute.These fixes are present in every rung.Source: §3.3 · p. 11

4 End-to-End RL Optimization

FlashNext-RL integrates rollout collection, training, weight transfer, and checkpointing across one training node and three inference nodes. Its memory controls, transfer path, and checkpoint optimizations address the constraints of running long-horizon RL within the training node’s host-memory budget.

4.1 System Overview

  • The final system places the trainer on one 8×B200 node and uses three inference nodes, with 672 rollouts in flight and a 131,072-token context cap.Generation overlaps training; after each step, rank 0 broadcasts updated weights to 24 inference GPUs.Source: §4.1 · p. 13, Figure 2 · p. 14

4.2 Host Memory

  • The training node has about 2.0 TB of host memory, while the final trainer configuration holds about 1.4 TiB of offloaded state plus the shared 95 GiB n-gram table.The system tracks proportional set size (PSS), because summed RSS overstates trainer memory by 43% when ranks map the shared table.Source: §4.2 · p. 14, Table 13 · p. 14
  • The node-local harness tarball reduces sandbox setup from 17.9 to 4.7 s and environment-worker memory from 200 to 17.6 GB in a 20-sandbox test.Minimum available memory rises from 82 to 289 GiB in a short test run.Source: §4.2 · p. 15
  • Exact-size arenas keep trainer PSS near pageable-offload levels while the StreamingMuon step falls from 78 to 28 to 12.5 s and the full step from 415 to 330 to 216 s.FSDP2’s default pinned offload was killed at 1,887 GiB; freezing the first two layers saves another 45 GiB of host memory.Source: §4.2 · pp. 15–16, Figure 3 · p. 16
  • For a 39M-token batch, streaming encode, zero-copy packing and per-micro-batch decoding reduce repeated batch holding; per-rank decode peak falls from 14.7 to 5.3 GiB.Per-CP payloads also reduce the receive-time memory dip from 104–108 GiB to 31–33 GiB on a 16M-token test batch.Source: §4.2 · pp. 15–16

4.3 Weight Transfer

  • Weight updates take 15–18 s from the trainer’s first send until engines resume, and the live-run broadcast averages 18.6 s.The transfer path gathers and converts weights layer by layer, omits the n-gram table, bounds messages to at most 1 GiB, and reloads layers on each engine.Source: §4.3 · p. 17

4.4 Checkpointing

  • DCP checkpoints go to local NVMe, while per-tensor arena storage avoids copying about 940 GiB of parameter shards and moments during a save.The frozen n-gram table is excluded from DCP save and load because it never changes and resides in shared memory.Source: §4.4 · p. 18
  • Overlapping HF export with broadcast cuts the checkpoint step from 307 s to 102 s in the real run.The streamed export uses four pinned buffers; 20 of 49 layer files come from the broadcast, and the background write finishes 132.5 s after starting.Source: §4.4 · p. 18, Table 15 · p. 18

5 A Real SWE RL Run

The live SWE RL run uses a target batch of 1,024 trainable rollouts with one training node and three inference nodes. Across steps 5–15, iterations average 22.4 minutes, and the reported entropy and trainer–sampler mismatch KL remain within bounded ranges.

Setup.

  • The run targets 1,024 trainable rollouts per batch and collects 1,249–1,584 rollouts per step, training 1,024–1,031 of them.Groups with all-zero advantages carry no learning signal and are not used for training.Source: Setup · p. 19
  • 1,344 s (22.4 min) is the average iteration time across steps 5–15, with 1,013 s of trainer work and 331 s waiting for rollouts.The trainer processes about 36.5M packed tokens during forward and backward passes; the loop is bound by generation.Source: Setup · p. 19, Table 16 · p. 19
  • The median rollout is 29K tokens, the p90 is 61K, and the p99 is 108K; rollouts longer than 64K are 7.7% of rollouts but contain 19% of tokens.Some trajectories reach the 131,072-token cap.Source: Setup · p. 19, Figure 4 · p. 20
  • After step 1, entropy stays within 0.66–0.76 nats per token and trainer–sampler mismatch KL with router replay stays within 3.6–4.0×10^-3.These measurements show no numerical collapse over the observed run, but do not establish long-run convergence or task-quality improvement.Source: Setup · p. 19, Figure 5 · p. 20
  • The estimated 8–11 training nodes, or 11–14 nodes including inference, are illustrative requirements under ideal scaling, not measured resource needs.The estimate applies slowdown proxies to trainer work; overlap with generation could mean less training capacity is needed to match the 22.4-minute iteration time.Source: Setup · pp. 20–21

6 Conclusion

FlashNext-RL combines a single-node BF16 trainer with vLLM rollout generation to optimize long-horizon RL for Qwen3.8-Flash-Next. The live SWE run supports a target batch of 1,024 trainable rollouts on one training node and three inference nodes, while the evaluation remains limited to a short run and this hardware configuration.

Conclusion

  • The study optimizes long-horizon agent RL for Qwen3.8-Flash-Next with a single-node BF16 trainer and vLLM-based rollout generation.The evaluation is limited to this hardware configuration and a short training run; it does not establish long-run convergence, final task accuracy, or multi-node training scalability.Source: §6 · p. 21
  • The live SWE system supports a target batch of 1,024 trainable rollouts with one training node and three inference nodes.The study concludes that adapting the offload path, parallel layout, and backward kernels to the model architecture is important for efficient RL training.Source: §6 · p. 21

Related papers

Loading 2610.28998v1…