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

Zhewei Yao; Yuxiong He

cs.LGOct 20, 2026

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

Authors: Zhewei Yao, Yuxiong He

Organizations: Snowflake AI Research

Abstract

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.

Figures & tables

Sessions/GPU	TP2 + EP (KV 1.76M)	TP4 + EP (KV 5.95M)	TP8 + EP (KV 8.41M)
7	835	715	–
10.5	985	814	–
17.5	840 (prefix hit 0.06)	943	658
26	771 (318 preemptions)	1,091	567
35	783 (preempting)	788	530

Table 1: BF16 engine on B200 GPUs uses replayed SWE rollouts and MTP3. Output tokens/s per GPU; “KV” is the KV-cache capacity per engine in tokens. The best measured configuration per row is in bold.

Optimizer	Bytes/param	Host for parameters	+ n-gram table ( $\sim$ 95.4 GiB)
AdamW	16	1,922 GiB	2,018 GiB (does not fit)
Muon budget	12	1,442 GiB	1,537 GiB (fits)

Table 2: Approximate host-memory budget for 129B parameters before layer freezing, on a node with about 1,937 GiB available. The Muon row assumes one fp32 moment per parameter for budgeting; the implementation uses AdamW for the remaining parameter groups. Final measured state sizes are in Table 13 .

	CP 1, 32K		CP 4, 131K
Offload	Muon step	Step	Muon step	Step	REAL tok/s/GPU
Pageable	$\sim$ 75 s	–	66 s	349 s	$\sim$ 426
PinnedArena	20–28 s	161–179 s	28 s	288 s	$\sim$ 516
+ pinned gradients	12.4 s	121–131 s	12.5 s	242 s	571–638

Table 3: Host offload path, full model on one node. Left: CP 1 with 32K bins (about 0.64M tokens per step). Right: CP 4 with a 131K bin.

Gradient path	Micro step	Step	GPU peak (alloc.)	REAL tok/s/GPU
FSDP2 host accumulation	$\sim$ 51 s	242 / 243 s	118.5 GiB	571 / 589
GPU-resident expert gradients	13–16 s	101 / 91 s	149.3 GiB	1,368 / 1,571

Table 4: GPU-resident expert gradients, full model, CP 4, 131K bin, 16-sample batches (about 1.1M tokens per step), steps 2 and 3.

Tokens/rank	DeepEP	MoK	Speedup
16K	32.2 ms	20.7 ms	1.56 $\times$
64K	130.1 ms	67.6 ms	1.92 $\times$
131K	294.0 ms	131.5 ms	2.24 $\times$

Table 5: MoE layer at EP 8. Left: forward + backward, DeepEP with grouped GEMMs vs. MoK (upstream build: intermediate size padded to 768, 24/28 communication SMs). Right: MoK at 32K tokens per rank, forward + recompute + backward.

Sequence length	PyTorch	cuDNN fused
8K	1.7 ms	0.2 ms
64K	35.3 ms	6.1 ms
131K	123.8 ms	22.5 ms

Table 6: QSA indexer + top-k, one B200, real layer inputs.

Implementation	16K	64K	dK/dV kernel at 64K
cuDNN DSA (MLA kernels, padded heads)	53.7 ms	219.8 ms	–
GQA-native Triton	33.9 ms	164.5 ms	97 ms
+ TMA gathers in dK/dV	29.0 ms	138.5 ms	76.2 ms
+ CuTe-DSL dK/dV (final)	20.7 ms	106.4 ms	46.7 ms

Table 7: QSA attention, forward + backward per layer, real layer-7 selections, one sequence (one B200).

	Eager	torch.compile	Fused
16K	9.55 ms	4.40 ms	3.42 ms
131K	75.17 ms	34.74 ms	17.84 ms
Kept at 131K	15.79 GiB	5.79 GiB	0.71 GiB

Table 8: One gated-residual read + inject, forward + backward, one B200. “Kept” is the memory saved for the backward at 131K.

$T$	CPU, 1 thread	CPU, 16 threads	UVA on GPU
16K	40–42 ms	14 ms	5.2 ms
131K	395–464 ms	117–131 ms	42 ms (no host cores)

Table 9: N-gram row gather for

T

tokens (

[T,16\times 160]

BF16 rows from the frozen table), one B200.

Figure 1: Cumulative trainer optimizations of FlashNext-RL on the 8-layer model: REAL tokens/s/GPU (top) and TFLOPs/GPU (bottom). Real rollouts, one node, 131K bins.

Removed	REAL tok/s/GPU	Change
MoK, back to DeepEP	13,699	$-$ 22.4%
Fused gated residual and glue	14,191	$-$ 19.6%
GQA-native QSA, back to cuDNN DSA	15,028	$-$ 14.9%
UVA n-gram gather, back to CPU	16,058	$-$ 9.0%
MoK communication SMs, back to 24/28	16,845	$-$ 4.6%
CuTe dK/dV, back to Triton TMA	16,966	$-$ 3.9%

Table 11: Removing one optimization at a time from the final 8-layer stack (L14; 17,651 REAL tokens/s/GPU on the same node).

Configuration	REAL tok/s/GPU	HW TFLOPs/GPU
First full-model step (CP 1, 32K bins, pageable offload)	$\sim$ 395	20
CP 4, 131K bin, pageable offload †	$\sim$ 426	21
+ PinnedArena †	$\sim$ 516	26
+ GPU-resident expert gradients †	1,368–1,571	68–78
+ GQA-native Triton QSA (real steps from here on)	2,297–2,333	114–116
+ MoK, static placement, pre-warm, GDN transpose drop	2,878–3,090	143–154

Table 12: Full 48-layer RL trainer development on one node. The initial row uses CP 1 and 32K bins; subsequent rows use CP 4 and 131K bins. Rows marked † use 16-sample batches (about 1.1M tokens); later rows use the earlier rollout batches of 7.1–7.6M tokens (median sample 66K). HW TFLOPs count matmuls only (49.8 GFLOP per token before freezing, 48.1 after, on these steps). Configurations and batch sizes vary, so the endpoints are not a controlled speedup comparison.

Panel of Figure 2 — Figure 2: The RL setup of FlashNext-RL: one training node and three inference nodes, with the orchestrator on the first inference node. Blue: NCCL weight broadcast; orange: rollouts and training batches (with routed experts); gray: checkpoints.

Structure	Size per rank	Where
fp32 parameter shards (incl. the 2 frozen layers)	58.9 GiB	PinnedArena
StreamingMuon moments	56.8 GiB	PinnedArena
fp32 gradients	56.1 GiB	gradient arena
n-gram table (BF16)	95.4 GiB per node	shared tmpfs mapping

Table 13: Host memory per trainer rank (8 ranks on the training node), final configuration.

Optimization	Problem	Saved
PinnedArena (Section 3.2.1 )	Caching host allocator rounds up to $2^{k}$ ; in-place registration crashes	fits (default pinned offload killed at 1,887 GiB)
Gradient arena	Pinned gradients rounded to $2^{k}$ and never freed	$\sim$ 105 GiB per node
Re-freeze after EP	Frozen layers’ experts got gradients and moments	37 GiB per node
Lazy optimizer state	450 GiB allocated during the sandbox startup burst	450 GiB off the startup peak
Batch copies	Two batches at the step boundary; int32 routes on the host; orchestrator copies	$\sim$ 300 GB (int32 routes) + one batch + $\sim$ 37 GB per buffered batch
CP host shard	Full routed experts on every CP rank	$\sim$ 110 GB per node

Table 14: Memory optimizations on the training node.

Figure 3: The full model on one node as the trainer’s memory optimizations are added: trainer host memory (PSS), GPU memory and StreamingMuon step time. FSDP2’s default pinned offload was killed before finishing a step.

Exporter	DCP	HF weights (on path)	Checkpoint-step overhead
Streamed per layer	71 s	236 s	307 s
From the broadcast, 4 buffers	70 s	32.3 s ∗	102 s

Table 15: Checkpoint cost on the training path. ∗ The background write finished 132.5 s after it started.

Figure 4: Rollout lengths in the real SWE RL run. Left: distribution over steps 5–15 (15,323 collected rollouts), with its median, p90 and p99. Right: median, mean and p90 per step (step 2’s traces were not saved).

Figure 5: Entropy and trainer–sampler mismatch KL per step in the real SWE RL run.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,742 papers

QwenGyre: An Elastic Reinforcement Learning Framework for Training xLong-Horizon Agents
Sep 27, 2026Weiqi Wang, Yuxin Zhou, Mouxiang Chen +9Reinforcement Learning Long-Horizon Agent Tasks
Alibaba Token Hub, Alibaba Group · University of Science and Technology of China · Tsinghua University
QaRL: Rollout-Aligned Quantization-Aware RL for Fast and Stable Training under Training--Inference Mismatch
Apr 9, 2026Hao Gu, Hao Wang, Jiacheng Liu +9Reinforcement Learning RL Fine-Tuning
The Hong Kong University of Science and Technology · City University of Hong Kong · Zhejiang University
LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget
Jul 16, 2026Changhai Zhou, Kieran Liu, Yuhua Zhou +17Efficient Language Model Training Long-Context Modeling
MindLab · Fudan University

Sep 27, 2026J/K move · Enter open · S save

QwenGyre: An Elastic Reinforcement Learning Framework for Training xLong-Horizon Agents

Weiqi Wang, Yuxin Zhou, Mouxiang Chen, Siyuan Zhang+8

Alibaba Token Hub, Alibaba Group · University of Science and Technology of China · Tsinghua University

Reinforcement Learning Long-Horizon Agent Tasks Long-Horizon LLM Agents RL for LLM Agents Long-Horizon RL GPU-Accelerated RL

Open paper AskPDF arXiv

Large language model (LLM) agents increasingly undertake extreme-long (xlong) horizon tasks, where a single execution can span hours, hundreds of model--environment interactions, and nearly 1M tokens per rollout. Applying online reinforcement learning (RL) to such executions poses two fundamental challenges: (1) severe execution variance and prolonged rollout delays cause massive GPU idling; and (2) complex non-linear branching generates massive trajectory redundancy, crippling training efficiency. To address these, we presents QwenGyre, an end-to-end framework for xlong-horizon online RL. QwenGyre elastically reallocates GPUs between rollout and training without interrupting live executions, while its trajectory processor reconstructs branching histories, scores partial progress, and deduplicates redundant paths to bound training costs. Scaled to our flagship model, Qwen~3.8 2.4T, with 700K tokens per rollout, QwenGyre yields a 6.0% absolute gain on NL2RepoBench (52.5%

\to

58.5%) in 48 steps. Across our evaluations on diverse domains of training datasets, QwenGyre delivers up to

1.85\times

and

1.78\times

speedups over Colocate and Async, respectively.

Abstract

Reinforcement Learning

Explore similar work

QwenGyre: An Elastic Reinforcement Learning Framework for Training xLong-Horizon Agents

QaRL: Rollout-Aligned Quantization-Aware RL for Fast and Stable Training under Training--Inference Mismatch

LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget