1 Introduction↩︎

Modern large language models (LLMs) [1]–[3] have become increasingly central to interactive, agentic, and embodied AI systems, spanning general-purpose assistants, personal agents [4], embodied AI, robotics control, and autonomous driving [5]–[8]. As these models move toward end-user, edge-device, and real-time service deployment, inference must preserve model capability while satisfying system constraints on latency, throughput, memory footprint, and energy consumption [9]–[11]. These constraints are especially pronounced in latency-sensitive applications and directly affect deployability on personal devices, edge hardware, and closed-loop decision-making systems. Consequently, serving efficiency and inference latency have become central bottlenecks for practical AR LLM deployment.

Efforts to improve LLM serving efficiency have proceeded along two complementary directions: making each model invocation cheaper, and reducing the number of serial invocations. The first direction relies on architectural changes. Compressed or sparse variants of softmax attention [12]–[15] and linear-attention or related recurrent-memory modules [16]–[20] reduce the storage and computation required by long-context and incremental decoding. Since purely linear memory can trade off retrieval fidelity, recent efficient LLMs increasingly use hybrid-attention backbones that combine softmax and linear-attention layers, balancing model capability, cache cost, and serving efficiency [21]–[27]. Hybrid attention lowers per-forward cost, but does not remove the serial nature of AR generation. The second direction reduces serial decoding steps. Multi-token prediction and speculative decoding retain the AR factorization while advancing multiple tokens per round through auxiliary heads, drafting, or parallel verification [28]–[32]. Diffusion large language models (dLLMs) go further by relaxing strict next-token generation and using iterative parallel denoising, offering a distinct route to lower decoding latency [33]–[37]. As illustrated in Figure 2, these trends have largely evolved in parallel in the open-source LLM ecosystem, with hybrid-attention AR models improving architectural efficiency and dLLMs improving generation-side parallelism.

Research on dLLMs has advanced diffusion-style objectives, masked or uniform diffusion, block diffusion, and large-scale diffusion LMs [38]–[45], with recent studies beginning to characterize their scaling and data-efficiency behavior [46]–[48]. In parallel, AR-initialized transfer and related conversion recipes obtain parallel-generation capability from existing AR checkpoints through continued training, shifted-token losses, block-diffusion training, or lightweight post-training [49]–[54]. Nevertheless, AR-to-dLLM transfer remains at an early stage: the roles of transfer-stage data, scaling behavior, and joint loss–mask design in capability preservation remain poorly understood. As a result, existing conversion methods often struggle to robustly inherit the capabilities of their AR seed checkpoints. At the same time, higher tokens-per-forward does not automatically translate into higher tokens-per-second; realized serving throughput depends on inference-system support for prefix reuse, KV or recurrent-state caching, and low-overhead parallel decoding [37], [50], [55]. On modern hybrid-attention backbones, these issues become more complex because diffusion training and inference must also handle recurrent-state handoff, mask realization, and serving-path design. How to convert a strong AR model into a capable and serving-efficient hybrid-backbone dLLM under a limited training budget therefore remains an open problem.

Motivated by these challenges, we present FLARE, a systematic recipe for converting hybrid-attention AR LLMs into capable and serving-efficient dLLMs under a practical training budget. FLARE examines the full conversion pipeline, from transfer-data construction and loss–mask design to hybrid-backbone training and serving-time inference. Concretely, FLARE makes three contributions. (1) Transfer diagnosis. To understand the performance degradation observed in prior AR-to-dLLM methods, we conduct controlled ablations over transfer data, objective design, and attention-mask choices, and identify transfer-data quality and distribution match as the dominant factors for preserving AR capability, outweighing loss formulation and mask design. This addresses a factor that has often been under-discussed in prior transfer work [50]–[52], [56]. (2) Efficient training. To enable diffusion training on modern softmax-plus-linear-attention backbones, we develop hardware-aware algorithms for computing linear attention under diffusion-specific visibility patterns [19], [20], [22], realizing these patterns through recurrent-state handoff rather than dense masking. (3) Efficient generation. To turn parallel-generation capability into wall-clock throughput, we build a unified inference system that supports both AR-style verified decoding and diffusion-style parallel denoising from the same checkpoint, extending efficient LLM serving ideas [9] to hybrid-backbone dLLMs.

Taken together, the study yields four main takeaways. 1) Objective and mask design matter for capability preservation: a token-balanced clean/noisy loss combines AR next-token supervision with block-diffusion denoising, while the document-packed clean/noisy mask enables both causal and bidirectional training signals without cross-document leakage. 2) Transfer data explains the remaining variation: once the core loss–mask design is aligned, transfer-data quality and distribution match dominate the residual performance gap, making AR-SFT a useful low-cost proxy for data selection. 3) Hybrid-backbone diffusion requires specialized training: diffusion visibility on linear-attention components cannot be imposed as a dense mask alone, and FLARE realizes it through recurrent-state scheduling and hardware-aware training kernels. 4) One checkpoint supports two generation regimes: starting from Qwen3.5 hybrid-attention checkpoints [27], FLARE-2B/4B/9B achieves competitive diffusion-LM quality while supporting both AR-style verified speculative decoding and diffusion-style parallel denoising (Figure 1).

2 Related Work↩︎

Within the AR factorization, multi-token prediction and speculative decoding relax the one-token-per-forward bottleneck by drafting several future tokens and verifying them with a target model [28]–[32]. These methods preserve the AR distribution under verification rules, but they still rely on a left-to-right reference model. Diffusion language models instead learn to fill masked positions through iterative denoising, allowing multiple token positions to be committed in one or a few forward passes. This line includes discrete diffusion models [38]–[41], [43], large-scale dLLMs such as LLaDA, Dream, Mercury, Seed Diffusion [33]–[35], [44], [45], [57], [58], and inference-oriented systems such as DFlash [37]. Block-diffusion LMs [42] provide a middle ground: blocks are generated sequentially, while positions inside a block can be denoised in parallel. A group of works reduces the training cost of this paradigm by converting or adapting pretrained AR checkpoints into dLLMs [49]–[54], [56]. These works make AR-to-diffusion transfer practical, but the sources of transfer degradation remain entangled across objective design, attention masks, transfer data, and evaluation protocol.

A separate line of work reduces the cost of each forward pass. Linear-attention, state-space, and recurrent-memory models replace the growing softmax-attention KV cache with a compact state [16]–[19], [59]–[63]. Delta-rule and gated linear-attention variants further improve the recurrent update used by modern efficient LLMs [20], [22], [64]. Related test-time-state models such as TTT and Atlas also fit the broader recurrent-memory trend shown in Figure 2 [65]–[67]. This improves memory scaling, especially for long contexts, but pure linear stacks can still be weaker than softmax attention on some tasks. As a result, recent efficient LLMs increasingly use hybrid backbones that interleave softmax and linear-attention layers, including Samba, Hymba, Zamba, MiniMax, Qwen3-Next/Qwen3.5, Kimi Linear, and Nemotron 3 [21]–[27], [68]–[71]. These architectures reduce cache size and long-context cost while largely retaining AR training and decoding.

Taken together, the timeline in Figure 2 shows that open-source LLM efficiency has advanced along two largely separate threads: generation-side parallelism, including dLLMs, and architecture-side reductions in per-forward cost through hybrid attention. Yet these threads have rarely been joined: dLLM work mostly assumes full-attention backbones, while hybrid LLMs largely retain AR training and decoding. FLARE studies this missing intersection as AR-to-diffusion post-training for hybrid checkpoints, adding a diffusion-style generation path to the same parameters while preserving causal generation. This requires jointly specifying the objective, clean/noisy mask, transfer data, recurrent-state implementation, and decoding interfaces developed next.

3 Method↩︎

3.1 Preliminaries on Attention Mechanisms↩︎

To ground the subsequent analysis of hybrid-attention backbones and non-causal visibility patterns, we introduce softmax and linear attention under a unified associative-memory view [72], [73], which abstracts attention layers as modules that memorize and retrieve. For notation, we denote the per-token queries and keys \(\mathbf{q}_t,\mathbf{k}_t\in\mathbb{R}^{d_k}\) and values \(\mathbf{v}_t\in\mathbb{R}^{d_v}\), received by each layer. It emits an output \(\mathbf{o}_t\in\mathbb{R}^{d_v}\). Under this view, softmax and linear attention share a memory–retrieval interface, but differ in how past \((\mathbf{k}_i,\mathbf{v}_i)\) pairs are stored and queried by \(\mathbf{q}_t\).

Softmax attention as explicit key-value memory. A softmax layer explicitly stores all visible key-value (KV) pairs \((\mathbf{k}_i,\mathbf{v}_i)\) in a KV cache. Given a query \(\mathbf{q}_t\), retrieval is performed by normalizing query-key similarities over a visible index set and taking a weighted sum of the corresponding values. Using a generic visible set \(\mathcal{V}_t\), the softmax retrieval is \[\mathbf{o}_t = \sum_{i\in\mathcal{V}_t} \frac{\exp(\mathbf{q}_t^{\!\top}\mathbf{k}_i/\sqrt{d_k})}{\sum_{j\in\mathcal{V}_t}\exp(\mathbf{q}_t^{\!\top}\mathbf{k}_j/\sqrt{d_k})} \mathbf{v}_i, \label{eq:softmax-retrieval}\tag{1}\] where \(\mathcal{V}_t\) denotes the visible index set for position \(t\); for full attention, \(\mathcal{V}_t\) contains all positions, while for causal attention it contains positions no later than \(t\). This explicit key-value memory provides accurate token-level access and long-range retrieval capability, but its KV-cache footprint grows as \(\mathcal{O}(T(d_k+d_v))\); dense full-sequence attention costs \(\mathcal{O}(T^2)\), while incremental decoding still requires \(\mathcal{O}(T)\) retrieval work per generated token.

Linear attention as recurrent memory. Linear attention [16], [59] replaces the explicit KV cache with a bounded-size recurrent state \(\mathbf{S}_t\in\mathbb{R}^{d_k\times d_v}\). In its simplest form, the state accumulates past key-value information through outer-product updates, \[\mathbf{S}_t = \mathbf{S}_{t-1} + \mathbf{k}_t\mathbf{v}_t^{\!\top}, \qquad \mathbf{o}_t = \mathbf{S}_t^{\!\top}\mathbf{q}_t . \label{eq:linear-retrieval}\tag{2}\] Retrieval is therefore performed from this compressed state rather than through explicit attention over stored key-value pairs. Modern efficient LLMs often instantiate recurrent memory with gating, decay, or delta-rule updates, including Gated DeltaNet [20] and Kimi Delta Attention [22]. These mechanisms share a bounded-state interface but differ in their concrete parameterizations. In this work, we use Gated DeltaNet (GDN) as the representative instantiation for exposition. Its memorization and retrieval rule, the Gated Delta Rule (GDR), is \[\mathbf{S}_t = \alpha_t\,\bigl(\mathbf{I}-\beta_t\,\mathbf{k}_t\mathbf{k}_t^{\!\top}\bigr)\mathbf{S}_{t-1} + \beta_t\,\mathbf{k}_t\mathbf{v}_t^{\!\top}, \qquad \mathbf{o}_t = \mathbf{S}_t^{\!\top}\mathbf{q}_t / \sqrt{d_k}. \label{eq:gdn-recurrence}\tag{3}\] The state-scheduling principle developed later relies on this recurrent-state interface rather than the specific parameterization of GDN, and can extend to other linear-attention variants. Compared with softmax attention, recurrent-memory attention maintains an \(\mathcal{O}(d_kd_v)\) state independent of sequence length; processing a length-\(T\) sequence costs \(\mathcal{O}(T)\) time, and incremental decoding uses constant-size state updates and retrieval per generated token, at the cost of a bounded and lossy memory.

Hybrid attention backbones. The two mechanisms above occupy different points on the spectrum of retrieval fidelity, memory compression, and serving efficiency. Softmax attention preserves explicit key-value memory and supports accurate token-level and long-range retrieval, as in Eq. 1 , while recurrent-memory attention compresses history into bounded states, as in Eq. 3 , reducing the dependence of memory on context length at the cost of lossy retrieval. Hybrid-attention backbones combine these mechanisms at the architecture level: softmax layers provide high-fidelity access to selected context, whereas recurrent-memory layers act as efficient compressed memory. By controlling the ratio and placement of these layer types, recent efficient LLM families seek to preserve model capability while reducing cache size and long-context serving cost [21], [22], [25]–[27], [69], [70]. Our main FLARE instantiation follows this setting. A hybrid backbone contains both explicit key-value memory and recurrent-state memory, so the same non-causal visibility pattern must be implemented through different mechanisms across layer types; the next subsection instantiates this issue for the clean/noisy training mask.

3.2 Training Objective and Mask Design↩︎

Let \(\mathbf{x}=(x^{1},\ldots,x^{L})\in\mathcal{V}^{L}\) be a training sample partitioned into \(K\) contiguous blocks \(\mathbf{x}_{b}=(x^{(b-1)B+1},\ldots,x^{bB})\) of size \(B\) (\(L=KB\)); we write \(\mathbf{x}^{<b}\) for all tokens strictly before block \(b\). For each block, we sample a masked subset \(\mathcal{M}_{b}\subseteq\{(b-1)B+1,\ldots,bB\}\) and form the noisy block \(\tilde{\mathbf{x}}_{b}\) by replacing \(x^{\ell}\) with [MASK] for \(\ell\in\mathcal{M}_{b}\). The complement \(\mathcal{M}_{b}^{c}\) defines a second, disjoint noisy view \(\tilde{\mathbf{x}}_{b}^{c}\). Based on this block-level corruption process, FLARE maximizes a token-balanced clean/noisy objective that preserves the AR path while adding diffusion-style supervision. The clean stream follows the original causal order and provides the next-token prediction term \(\mathcal{L}_{\mathrm{AR}}\). The noisy stream provides the diffusion term \(\mathcal{L}_{\mathrm{diff}}\): the two complementary noisy views predict disjoint token subsets, both conditioned on the same preceding clean context. The overall objective is \[\mathcal{L}_{\boldsymbol{FLARE}}(\theta) = \mathcal{L}_{\mathrm{AR}}(\theta) + \mathcal{L}_{\mathrm{diff}}(\theta). \label{eq:flare-loss}\tag{4}\] Concretely, the clean-stream term is the usual next-token log-likelihood over the full sequence, and the noisy-stream term sums the denoising log-likelihood over the two complementary masked views: \[\begin{align} \mathcal{L}_{\mathrm{AR}}(\theta) &= \sum_{\ell=1}^{L} \log p_\theta\!\left(x^{\ell}\mid \mathbf{x}^{<\ell}\right), \nonumber \\ \mathcal{L}_{\mathrm{diff}}(\theta) &= \sum_{b=1}^{K} \left[ \underbrace{ \sum_{\ell\in\mathcal{M}_{b}} \log p_\theta\!\left(x^{\ell}\mid \tilde{\mathbf{x}}_{b},\mathbf{x}^{<b}\right) }_{\text{primary noisy view}} + \underbrace{ \sum_{\ell\in\mathcal{M}_{b}^{c}} \log p_\theta\!\left(x^{\ell}\mid \tilde{\mathbf{x}}_{b}^{c},\mathbf{x}^{<b}\right) }_{\text{complementary noisy view}} \right]. \label{eq:flare-loss-terms} \end{align}\tag{5}\]

The clean/noisy construction in Figure 3 realizes the two terms in Eq. 5 : the causal clean stream supplies \(\mathcal{L}_{\mathrm{AR}}\), while the block-bidirectional noisy stream supplies \(\mathcal{L}_{\mathrm{diff}}\) under two complementary views. This differs from standard mask-dLLM training [42], [44], [51], where masked positions are predicted only from noisy contexts. In FLARE, every token contributes one AR signal and one diffusion signal at unit weight because \(\mathcal{M}_{b}\) and \(\mathcal{M}_{b}^{c}\) partition each block. This keeps the two supervision sources balanced under varying mask fractions while giving the same checkpoint both causal generation and block-level parallel denoising capabilities. Under the memory view above, an attention mask specifies which token information is accessible when a query reads from memory. Figure 3 shows the document-packed clean/noisy visibility pattern used to compute the objective above in a packed forward pass: the clean stream remains token-causal and isolated from the noisy stream, each noisy block is bidirectional within the block and attends to preceding clean context, and document boundaries isolate packed samples to prevent cross-document leakage while avoiding padding overhead. The two memory implementations in a hybrid backbone realize this same pattern differently: in softmax layers, an additive mask \(\mathbf{M}\in\{0,-\infty\}^{T\times T}\) inside the softmax of Eq. 1 directly removes forbidden entries from the visible set \(\mathcal{V}_t\); in recurrent-memory layers, making a past pair \((\mathbf{k}_i,\mathbf{v}_i)\) inaccessible to a later query means preventing its information from entering or propagating through the relevant state \(\mathbf{S}_t\). Thus the same non-causal clean/noisy mask becomes a schedule of writes, state propagation, and state resets rather than a post-hoc masking matrix; the corresponding recurrent-state schedule is derived below.

Algorithmically, this construction is closest to TiDAR [52] and concurrent work I-DLM [54], but differs in two choices that matter for decoding. Unlike TiDAR, we apply logit shift to the noisy-stream diffusion terms to align them with AR semantics. Unlike I-DLM, we keep the noisy stream block-bidirectional and use random rather than fully masked noisy views during training, which preserves Diffusion-Trust decoding in addition to AR-style decoding.

The clean/noisy mask above can be applied as a visibility matrix in softmax layers, but must be translated into a state schedule for linear-attention layers. We use the generic form \(\mathbf{S}_\ell=\mathcal{T}_\ell\mathbf{S}_{\ell-1}+\mathcal{W}_\ell\) with per-position output \(\mathbf{o}_\ell=f(\mathbf{S}_\ell,\mathbf{q}_\ell)\), which covers Mamba-2 [19], Gated DeltaNet [20], and Kimi Delta Attention [22]. We instantiate the construction with GDN below, while the same state-scheduling idea extends to other members of this family by substituting the corresponding transition, write, and output maps. Let \(e(b)=bB\); we index clean positions by \(t\), noisy positions by \(\ell\), and write noisy-side quantities with a tilde. For clarity, the equations below are written for a single document segment; under document packing, the same schedule is applied independently within each segment, with states reset at document boundaries.

For a GDN layer, the clean stream follows the standard causal recurrence of Eq. 3 . For noisy block \(b\), the recurrent state is initialized from the preceding clean boundary state, \(\tilde{\mathbf{S}}_{e(b-1)}:=\mathbf{S}_{e(b-1)}\), and then updated only with noisy tokens from the same block: \[\left\{ \begin{align} \mathbf{S}_t &= \alpha_t\bigl(\mathbf{I}-\beta_t\mathbf{k}_t\mathbf{k}_t^{\!\top}\bigr)\mathbf{S}_{t-1} + \beta_t\mathbf{k}_t\mathbf{v}_t^{\!\top}, && \text{clean stream}, \\ \tilde{\mathbf{S}}_\ell &= \tilde{\alpha}_\ell\bigl(\mathbf{I}-\tilde{\beta}_\ell\tilde{\mathbf{k}}_\ell\tilde{\mathbf{k}}_\ell^{\!\top}\bigr) \tilde{\mathbf{S}}_{\ell-1} + \tilde{\beta}_\ell\tilde{\mathbf{k}}_\ell\tilde{\mathbf{v}}_\ell^{\!\top}, && \text{noisy block } b, \end{align} \right. \label{eq:flare-gdn-schedule}\tag{6}\] where \(\ell\in\{e(b-1)+1,\ldots,e(b)\}\). The clean branch uses the standard causal readout \(\mathbf{o}_t=\mathbf{S}_t^{\!\top}\mathbf{q}_t/\sqrt{d_k}\), while the noisy branch reads every position from the shared block-end state \(\tilde{\mathbf{o}}_\ell=\tilde{\mathbf{S}}_{e(b)}^{\!\top}\tilde{\mathbf{q}}_\ell/\sqrt{d_k}\). This block-end readout realizes bidirectional visibility within the noisy block, and resetting the noisy state to the corresponding clean boundary state prevents information flow across noisy blocks. These recurrences define the intended visibility pattern, but the direct sequential form is not used for training; the chunk-parallel realization and kernel implementation are discussed in Section 5.2 and Appendix 7.

3.3 Decoding Interfaces↩︎

The clean/noisy objective in Section 3.2 gives a trained FLARE checkpoint two usable output distributions under the same forward interface: clean-stream logits preserve causal AR prediction, while noisy-stream logits denoise a bidirectional block conditioned on clean context. This separation is important for transfer: the seed checkpoint’s capability is organized around next-token prediction, and a pure diffusion path would need to replace that causal distribution under limited post-training data. The clean stream therefore anchors the transferred AR capability and provides a verification reference, while the noisy stream adds block-level parallel denoising. This yields two decoding interfaces. In Diffusion-Trust, noisy-stream samples are the reference and are committed through block-denoising confidence. In AR-Trust, noisy-stream samples serve as drafts and clean-stream logits verify them left-to-right. Algorithms 4 and 5 give the operational view; below, we formalize the corresponding commitment and verification rules.

We use \(x\) for the input prompt, \(y\) for the generated sequence, \(f_\theta\) for the trained checkpoint’s forward map, \(z\) for pre-softmax logits, and \(\sigma_{k,p,T}\) for temperature-scaled top-\(k\)/top-\(p\) sampling. Diffusion-Trust uses noisy-stream distributions \(\pi_i\); AR-Trust uses clean-stream target distributions \(p_i\), noisy proposal laws \(q_i\), and draft tokens \(d\). Decoding schematics and path-specific masks are provided in Appendix 8, while system-level serving details are discussed in Section 5.3.

We next state the two token-emission rules called by these interfaces.

Diffusion-Trust: block iterative denoising and noisy-stream commitment. This rule formalizes the commit operation invoked by the parallel denoising loop in Algorithm 4. Consider an active block \(\mathcal{B}_{t,B}=\{t,\ldots,t+B-1\}\) behind the causal prefix, decoded over steps \(s=1,\ldots,S\), with \(R_1=\mathcal{B}_{t,B}\). Let \(R_s\) be the unresolved positions before step \(s\). For each \(i\in R_s\), the noisy stream defines \(\pi_i=\sigma_{k,p,T}(\tilde{z}^{\mathrm{shift}}_i)\), samples \(\hat{y}_i\sim\pi_i\), and computes a confidence score \(c_i=\operatorname{score}(\pi_i,\hat{y}_i)\). The commit set \(A_s\) is selected in parallel across the unresolved positions, and since the noisy stream is trusted in this path, selected candidates are written to \(y\) without clean-stream verification: \[\tag{7} \begin{align} A_s &= \begin{cases} \{i\in R_s:c_i\ge\gamma_s\}, & s<S,\\ R_s, & s=S, \end{cases} \tag{8}\\ y_i^{(s+1)} &= \begin{cases} \hat{y}_i, & i\in A_s,\\ y_i^{(s)}, & i\in R_s\setminus A_s, \end{cases} \qquad R_{s+1}=R_s\setminus A_s . \tag{9} \end{align}\] Positions outside \(A_s\) stay masked for the next denoising step, and the case \(s=S\) commits all unresolved positions so the block is finalized after at most \(S\) steps. Eq. 7 covers confidence-based [44] and margin-based [74] unmasking by changing how \(c_i\) is computed. Because intermediate denoising forwards contain provisional block contents, their recurrent states are not written into the live causal state. Once the block is finalized, the causal replay step in Algorithm 4 commits the finalized tokens to the hybrid backbone state; the full serving loop is given in Appendix 8.3.

Eqs. 7 and 10 define the two regimes used by Algorithms 4 and 5. In AR-Trust, Eq. 10 accepts each held draft with probability \(\min\{1,p_i(d_i)/q_i(d_i)\}\); following the acceptance-rate analysis of I-DLM [54], the plotted \(\alpha\) abstracts the average success probability across horizon \(K\), so speed improves only when noisy proposals align with the clean target. In Diffusion-Trust, Eq. 7 commits an active block over at most \(S\) denoising steps, so finalizing \(B\) tokens in \(S\) forwards yields the ideal ratio \(B/S\). Figure 6 summarizes these analytical regimes and isolates the algorithmic control variables: \(\alpha\) and \(K\) for AR-Trust, and \(B\) and \(S\) for Diffusion-Trust. System factors such as state replay, cache updates, and kernel scheduling are deliberately excluded here; measured throughput is reported in Section 5.3.

4 Experimental Analysis↩︎

We use this section to identify the key algorithmic and data ingredients of AR-to-dLLM transfer. While existing conversion methods [50], [51], [56] report a range of transfer outcomes, several design choices — the training objective, the clean/noisy visibility pattern, whether the clean stream is explicitly aligned with AR next-token prediction, logit shift, transfer-data composition, and evaluation protocol — have not yet been examined systematically and in isolation. We therefore start from a smaller model, Qwen3-1.7B, and conduct a controlled study under a fixed training budget and evaluation suite, varying algorithmic ingredients and transfer data separately, so as to surface the recipe components that drive successful transfer. The resulting recipe is then carried over to stronger hybrid-attention models in Section 5.

All transfer runs initialize from the same pretrained Qwen3-1.7B checkpoint [75]. In our SFT format, the prompt tokens act as the leading clean-stream context (visible through attention but excluded from the loss), while the assistant response tokens are partitioned into blocks of size \(B\) that simultaneously receive the \(\mathcal{L}_{\mathrm{AR}}\) and \(\mathcal{L}_{\mathrm{diff}}\) terms defined in Section 3.2 — i.e., the clean/noisy dual-stream supervision is applied only on response tokens. All runs use a maximum sequence length of 4096 with the document-level packing mask of Figure 3, a global batch size of 256, and 9000 optimizer steps, corresponding to roughly 10B training tokens. We evaluate all checkpoints on the same 12-task suite, covering Math + Reasoning (GSM8K [76], MATH-500 [77], AIME 24/25 [78], ARC-Challenge [79], GPQA-Diamond [80]), Knowledge + IF (MMLU [81], MMLU-Pro [82], IFEval [83]), and Code (HumanEval [84], MBPP [85], LiveCodeBench v6 [86]). All checkpoints share one decoding protocol; detailed settings are in Appendix 10.

4.1 Effects of Transfer Data Composition↩︎

We first study transfer data because it determines the regime in which the algorithmic ablations should be interpreted. We construct four controlled transfer-data mixes from curated open-source post-training SFT corpora released with the Llama-Nemotron [87], Nemotron Nano 2 [23], Nemotron 3  [24], and Nemotron-Cascade-2  [88] model families; the detailed curation, filtering, packing, and mixture-sampling procedure is provided in Appendix 9. The mixes vary three factors while keeping the seed checkpoint, training budget, and evaluation protocol fixed: reasoning-trace length, math/domain coverage, and instruction-following coverage.

Concretely, the four mixes are built from three reusable source pools, each drawn from publicly released SFT corpora: a long chain-of-thought pool (Long-CoT, primarily Llama-Nemotron-Post-Training [87]), a math pool (Math, the math splits of Llama-Nemotron-Post-Training plus Nemotron-Math-Proofs-v1), and an instruction-following pool (IF, the IF-tagged subsets of Nemotron-Cascade-1/-2-SFT-Data and Nemotron-Instruction-Following-Chat-v1/v2 [23]); a short-reasoning pool (Short-CoT, the English split of Nemotron-Post-Training-Dataset-v2 [23]) is used only by Mix 2. Specifically, (1) Mix 1 (Long-CoT) uses the Long-CoT pool alone and serves as a reasoning-heavy baseline; (2) Mix 2 (Short-CoT+Math) combines the Short-CoT pool with the Math pool to separate the effect of trace length from math coverage; (3) Mix 3 (Long-CoT+Math) replaces Short-CoT with the Long-CoT pool to test whether stronger math coverage helps without shortening traces; and (4) Mix 4 (Long-CoT+Math+IF) further adds the IF pool on top of Long-CoT and Math to test whether broader instruction coverage improves transfer beyond reasoning-centric domains. The four pools span roughly \(0.5\)–\(17\)B assistant tokens, and components are combined under per-mix sampling weights that we arrived at empirically through the comparisons reported below; the resulting weights are visualized in Figure 7 (a) and detailed in Table 7 and Appendix 9.3, while Appendix 9.5 additionally documents an automatic, instance-level data-selection pipeline that we explored as a complementary route.

Figure 7 (b)–(d) compares these four mixes across four transfer recipes under the same Qwen3-1.7B seed and training budget. AR fine-tuning (AR-SFT) denotes continued next-token fine-tuning on response tokens; all-mask causal (AM-Causal) and all-mask bidirectional (AM-Bidir) share the joint AR/diffusion supervision used by FLARE(formally introduced in Section 4.2) but replace random noisy masking with all-masked noisy views, differing only in whether the noisy stream is token-causal or bidirectional within each block (AM-Causal is close in spirit to concurrent I-DLM [54]). We read the resulting comparison along two axes. First, on the recipe axis, the all-mask variants tend to trail both AR fine-tuning and FLARE—most visibly on Code, where they drop sharply on the math-heavy Mix 2 and Mix 4—while AR fine-tuning and FLARE track each other closely within each mix. The first observation flags an objective-side gap that we dissect under a fixed data mix in Section 4.2; the second indicates that, with a properly aligned clean/noisy objective, AR fine-tuning is a faithful low-cost proxy for screening data mixes before running more expensive dLLM conversion. Second, on the data axis, no single data change is uniformly beneficial: Mix 2 slightly improves Math + Reasoning but substantially weakens Code, whereas Mix 4 gives the strongest Knowledge + Instruction Following results while keeping Math + Reasoning and Code competitive.

Two conclusions follow from the recipe and data axes of Figure 7. First, conditional on a properly aligned clean/noisy objective, the attainable transfer quality is governed primarily by the transfer data mix rather than by the algorithmic recipe; in particular, converted dLLMs track their AR fine-tuning counterparts under the same mix up to capability-group-specific deviations, so AR fine-tuning provides a low-cost proxy for screening data mixes prior to dLLM conversion. Second, comparisons of objective and attention-mask choices across different mixes are confounded by data quality and are therefore uninformative as algorithmic ablations. These observations motivate the controlled design used in the remainder of Section 4: Section 4.2 fixes Mix 1 (Long-CoT) and isolates the effects of objective design, clean/noisy visibility, logit shift, and noisy-mask sampling, while Section 5 adopts Mix 4 (Long-CoT+Math+IF) for the final hybrid-checkpoint conversions on the basis of its balanced behavior across capability groups in this sweep.

4.2 Objective and Attention-Mask Ablations↩︎

After isolating the data effect in Section 4.1, we fix the transfer data condition to Mix 1 (Long-CoT) and ask which objective and attention-mask ingredients prevent AR-to-dLLM conversion from losing the seed model’s capability. All runs use the same document-level packing mask, which keeps variable-length prompt-response samples isolated inside packed batches while avoiding padding overhead. We therefore treat document-level masking as a fixed efficiency condition rather than an ablated factor, and focus on two conversion-specific choices: whether the clean stream retains AR-style token-causal supervision, and whether the noisy stream learns block-diffusion denoising.

Figure 8 (a) studies the clean-stream side through a cumulative recipe. Starting from AR fine-tuning, we replace next-token fine-tuning with a pure block-diffusion objective following the spirit of BD3-LM [42] and SDAR [51], then add back a token-causal AR clean stream, an auxiliary next-token-prediction loss on clean-stream logits, and the logit-shifted noisy-stream supervision used by FLARE. The results show that pure block-diffusion transfer alone is not capability-preserving: it substantially degrades all three capability groups (averaging \(-21.8\) points relative to AR fine-tuning, with the heaviest drops on Code and Knowledge + IF). The largest single-step recovery comes from restoring a token-causal AR clean stream (\(+14.0\) points on average), which alone closes most of the Code and Knowledge + IF gap. Adding the clean-stream next-token-prediction loss further pulls Math + Reasoning back to the AR fine-tuning level by anchoring the converted model to the pretrained checkpoint’s AR semantics, and logit shift on top yields a near-saturated final score.

Figure 8 (b) isolates the noisy-stream side. At fixed AR-aligned clean-stream supervision, all-masked noisy views and randomly masked diffusion views yield comparable category-level scores, so noisy-mask sampling is not a primary driver of benchmark recovery; its role is instead to determine decoding compatibility, since random masking trains the noisy stream as a genuine denoising path and preserves Diffusion-Trust decoding, while logit shift aligns the block-diffusion logits with inference-time token positions and avoids a wasted block-boundary prediction.

Two conclusions follow from the clean- and noisy-stream axes of Figure 8. First, capability-preserving AR-to-dLLM conversion is governed primarily by clean-stream alignment, with logit shift and noisy-mask sampling contributing only marginal additional benchmark accuracy. Second, these auxiliary ingredients are retained for decoding compatibility, since they preserve both AR-Trust and Diffusion-Trust sampling paths. Combined with Section 4.1, this isolates the residual transfer gap to the quality of the data mix.

5 Hybrid-Backbone Conversion and Evaluation↩︎

Section 4 decomposed AR-to-dLLM transfer into two complementary parts: an AR-aligned clean stream that determines whether capability is preserved, and a noisy-stream design that determines whether the converted checkpoint remains compatible with diffusion-style parallel decoding. This section applies the resulting configuration to strong hybrid-attention backbones and asks three concrete questions: (i) whether the transfer configuration selected by the controlled study scales to Qwen3.5 checkpoints [27] so that the converted models retain the capability of their AR source; (ii) whether two-stream supervision can be executed on a hybrid-attention backbone at hardware efficiency close to pure-AR training; and (iii) whether a single checkpoint can support both AR-Trust and Diffusion-Trust decoding under one unified serving stack. The three questions are addressed in turn in the subsections that follow.

We train FLARE-2B/4B/9B from the post-trained Qwen3.5-2B/4B/9B AR checkpoints, and compare the converted models against a set of leading dLLM systems: the commercial Mercury-2 [89], the LLaDA-2.0/2.1 mini and flash families [45], [57], and the SDAR family [51]. All three FLARE checkpoints are trained in a single supervised-finetuning stage on Mix 4 (Long-CoT+Math+IF), the data condition selected in Section 4.1. We keep the main protocol aligned with Section 4: sequence length \(4096\), document-level packing masks, global batch size \(256\), and \(9000\) optimizer steps, corresponding to roughly 10B training tokens. The objective is the token-balanced AR/diffusion loss \(\mathcal{L}_{\mathrm{FLARE}}\) in Eq. 4 , with block size \(B=4\). This conversion budget is markedly more modest than prior AR-to-dLLM transfer work, which typically uses 50B–200B tokens [50], [51]. Evaluation follows the 12-benchmark protocol from Section 4; we report FLARE under both AR-Trust and Diffusion-Trust sampling, with the original Qwen3.5 checkpoints serving as AR source-model references. The LLaDA and SDAR rows are taken from their technical reports, Mercury-2 is evaluated through its official API, and the Qwen3.5 and FLARE rows are measured under our SGLang-based inference stack [9].

5.1 Capability Retention Across Model Scales↩︎

We first evaluate whether the converted checkpoints preserve the capability of their AR initialization. Tables 1 and 2 report results at 9B and at smaller scales, with the original Qwen3.5 checkpoints serving as source-model references and the dLLM systems listed above as external baselines.

2pt 3pt 2.2pt

2pt 3pt 2.7pt

At the 9B scale, FLARE-9B matches or exceeds LLaDA-2.1-flash on most shared benchmarks, despite using roughly \(1/10\) of its total parameters: GPQA-Diamond (\(71.21\) vs.\(66.67\)), MMLU-Pro (\(77.39\) vs.\(75.31\)), MBPP (\(91.05\) vs.\(88.29\)), and LiveCodeBench v6 (\(49.71\) vs.\(44.05\)). Compared with the commercial Mercury-2 system, FLARE-9B is clearly stronger on math-reasoning benchmarks (MATH-500 by \(+14.2\) points, AIME-24 by \(+12.2\) points); we note in fairness a remaining gap on LiveCodeBench v6 (\(49.71\) vs.\(67.30\)), reflecting Mercury-2’s additional training advantage on code generation.

More directly, relative to the source-model Qwen3.5-9B, FLARE-9B retains the majority of the original capability across most tasks: \(98.6\%\) on MATH-500 (\(95.20\) vs.\(96.60\)), \(96.6\%\) on AIME-24 (\(63.33\) vs.\(65.56\)), and \(95.1\%\) on MMLU-Pro (\(77.39\) vs.\(81.39\)), while exceeding the original AR checkpoint on MBPP (\(91.05\) vs.\(89.11\)). The conversion thus preserves high-level reasoning and coding ability rather than simply improving easier short-form benchmarks. Together, the 9B results support both aspects of question (i) at the start of this section: parameter efficiency relative to substantially larger dLLM baselines, and capability retention relative to the same-family AR checkpoint.

The same pattern extends to the smaller scales. FLARE-4B outperforms SDAR-30B-A3B on nearly all reported benchmarks despite using a much smaller dense backbone, with the largest margins on GPQA-Diamond (\(63.64\) vs.\(36.7\)), MATH-500 (\(94.20\) vs.\(77.8\)), and AIME-24 (\(58.89\) vs.\(16.7\)). At the 2B scale, FLARE-2B improves substantially over the parameter-matched SDAR-1.7B on math and code tasks, including AIME-24 (\(31.11\) vs.\(10.0\)) and MATH-500 (\(84.40\) vs.\(63.2\)). Relative to the original Qwen3.5 checkpoints at both scales, the converted models retain a large fraction of the inherited capability, further indicating that the method functions as capability-preserving transfer rather than training a new dLLM from scratch.

Each FLARE row is reported under both AR-Trust and Diffusion-Trust sampling, but the two columns correspond to a single trained checkpoint that only switches its sampling path. The two modes are close on math and on most knowledge benchmarks (with larger AR-Trust advantages on GPQA-Diamond and IFEval at 9B), while Diffusion-Trust is consistently weaker on code-generation tasks. We attribute this gap to two factors: diffusion-style decoding is more sensitive to brittle answer extraction (which fails under minor syntax variations) and long-output truncation under the current code-evaluation scripts, but the parallel-decoding interface also faces an inherent challenge on long, strictly structured outputs such as multi-function code, where each block commits without left-to-right syntactic verification. Results marked with ^* should therefore be interpreted conservatively. Even with this caveat, the comparison supports the feasibility side of question (iii): a single FLARE checkpoint can support both a clean-stream verified causal sampling path and a block-diffusion sampling path, without training separate models for different serving regimes.

A residual gap to the original AR checkpoints remains, especially on instruction following and some coding benchmarks. We attribute this partly to the post-trained nature of the Qwen3.5 source models: continued SFT with external data drawn from a different distribution than the original post-training data can shift the output distribution away from the source checkpoint, consistent with the data-quality findings of Section 4. Future conversions may therefore benefit from transfer data more tightly matched to the source AR model’s own reasoning and instruction-following style. With capability retention established, the next subsection turns to question (ii): whether this two-stream supervision can be executed on a hybrid backbone at acceptable training efficiency.

5.2 Efficient Training on Hybrid Backbones↩︎

This subsection addresses question (ii): whether the transfer configuration can run on a hybrid-attention backbone at hardware efficiency close to pure-AR training. The clean/noisy mask in Section 3.2 defines a causal clean stream with block-bidirectional noisy streams. In softmax-attention layers this is just a visibility mask over key-value pairs, but in recurrent linear-attention layers visibility is encoded by the state trajectory itself: a noisy block must be initialized from the correct clean boundary state, expose its tokens to one another, and avoid leaking into unrelated blocks or packed documents. The same algorithmic mask thus becomes a state-scheduling problem inside the hybrid backbone.

On the Qwen3.5 backbone, this state-scheduling problem affects two sequence-mixing components in each GDN layer: the GDR recurrence and the width-\(W\) causal ShortConv that prepares its inputs, imposing three requirements absent in pure AR training, which sees one token-causal stream. (1) Gated Delta Rule: every noisy block must start from the clean state at its preceding block boundary; these mid-chunk states scale as \(L/B\) in a direct implementation and dominate memory at the small block sizes FLARE uses. (2) 1D causal ShortConv: noisy tokens must read both in-block noisy lags and clean-context lags from before the boundary, realized directly via \(L/B\) small block-level convolution launches. (3) Document packing: packed examples must stay isolated despite recurrent state propagation and convolutional lag reads.

We compare two realizations of this schedule, sharing clean/noisy visibility semantics but differing in how block-boundary states are produced and consumed. (1) Route I (chunk-then-refine) computes the clean stream first, materializes every block-boundary clean state \(\mathbf{S}_{(b-1)B}\) in HBM, then runs each noisy block as a seeded local recurrence; it serves as a correctness reference and reuses the standard AR kernel structure, but introduces \(L/B\)-scaling state hand-offs. (2) Route II (fused two-stream) stores only strided clean-state checkpoints, reconstructs the required boundary state in registers, immediately consumes the corresponding noisy block, and injects block-level gradients back into the clean recurrence without materializing intermediate boundary states; the ShortConv branch fuses lag selection, document-boundary masking, and noisy-block convolution into one kernel, using cu_seqlens-aware guards to reset noisy states, block cross-document gradient shifts, and zero ShortConv reads crossing document boundaries. Derivations and route-level pseudocode appear in Appendix 7.

We first evaluate the schedules at the kernel level. Figure 9 shows Route II is most effective in the small-block regime diffusion training requires. For the GDR at \(B{=}1\), it cuts total latency from \(135.10\) ms to \(37.69\) ms and peak memory from \(18.14\) GiB to \(0.45\) GiB. At larger \(B\), Route I can win because its dense chunk-level matmul better saturates tensor cores. Since diffusion training requires small block sizes (FLARE uses \(B=4\)), Route II is a structural requirement rather than an optimization: without it the small-\(B\) regime is infeasible in memory. We therefore use Route II for FLARE’s small diffusion blocks and retain Route I as the large-block path. For ShortConv, Route II stays faster across all measured block sizes by avoiding many small block-level launches.

Figure 10 shows these kernel-level gains carry to end-to-end training, measured as model FLOPs utilization (MFU). Replacing the GDR path with Route II removes the small-\(B\) memory failure and raises MFU at \(B{=}4\) from \(13.80\%\) to \(17.93\%\). The recovered memory headroom supports a larger per-GPU local batch, lifting MFU to \(21.40\%\). Adding the fused ShortConv path brings FLARE-2B to \(24.81\%\) MFU at \(B{=}4\), matching and slightly exceeding the pure-AR Qwen3.5-2B reference (\(24.04\%\)) on the same hardware. Since two-stream supervision processes more work per step than single-stream AR training, the cost of recurrent-state scheduling is fully absorbed by the hardware-aware implementation rather than becoming the dominant training bottleneck. This affirmatively answers question (ii); with training efficiency established, the next subsection turns to question (iii): how to support both decoding paths under one unified serving stack.

5.3 Unified Serving for AR-Trust and Diffusion-Trust↩︎

This subsection addresses question (iii): whether a single serving stack can support both decoding paths efficiently. Inference is the other stage where the hybrid backbone changes the standard dLLM serving problem. A pure softmax Transformer can treat speculative decoding largely as KV-cache management, while FLARE must coordinate a softmax KV cache, a GDN recurrent state, and the clean/noisy attention patterns used by the two decoding paths. We implement this in an SGLang-based serving stack [9] that runs the AR-Trust and Diffusion-Trust interfaces of Section 3.3 from the same checkpoint, with path-specific masks and shared hybrid-attention kernels. Full pseudocode for both paths is given in Appendix 8; the hybrid-specific serving machinery is detailed in Appendix 8.4.

Figure 6 gives the analytical interface-level trade-off between AR-Trust and Diffusion-Trust. Here we focus on the serving mechanisms that turn those regimes into efficient execution on a hybrid-attention backbone; measured tokens-per-second (TPS) is reported later in Figure 11.

The serving stack is built around four mechanisms beyond a standard speculative-decoding engine. (1) Recurrent-state commit: accepting only \(r\) of \(K\) speculative tokens is not just a KV tail trim, since the GDN state after the verify round must also be rewound to \(\mathbf{S}^{(r)}\). (2) Native dLLM mask modes: dense custom masks impose unnecessary per-score overhead, so FLARE encodes the clean/noisy row types compactly and uses a prefix-tile fast path whenever a KV tile lies entirely in the causal prefix. (3) Fused verification and top-\(k\) kernels: AR-Trust verification avoids materializing full \([M,V]\) logits or probability tensors, which is important at large vocabularies. (4) CUDA-graph replay safety: graph reuse must match not only tensor shapes but also block size, mask mode, recurrent-state update mode, and logits-output mode; otherwise a graph captured for one decoding path can be incorrectly replayed for another.

The most hybrid-specific of the four is mechanism (1). To make partial-accept rewind efficient, FLARE records the intermediate verify-position states inside the same recurrent kernel and commits the accepted offset with one fused gather-scatter kernel (Figure 31 in Appendix 8). This avoids replaying accepted tokens through every GDN layer after each verify decision. Diffusion-Trust uses a related separation between denoise and commit: denoise passes read the recurrent state but do not write it back, and a final causal state-update pass commits the completed block to the recurrent pool. The reason is that tokens in intermediate denoise rounds may still be revised in subsequent rounds, so writing them back early would contaminate the recurrent state trajectory used by later blocks.

Figure 11 reports the high-concurrency headline setting, where the system mechanisms above translate the parallel decoding advantage into real tokens-per-second. On a single A100-80GB at \(C{=}8\), FLARE-2B reaches \(2{,}087\) tokens/s on GSM8K, giving a \(2.2\times\) gain over LLaDA-2.1-mini and a \(4.8\times\) gain over SDAR-1.7B. On GPQA-Diamond, it reaches \(1{,}441\) tokens/s, a \(3.6\times\) gain over LLaDA-2.1-mini. The gap is largest at high concurrency, the regime in which per-step overhead dominates and the fused kernels are most valuable for keeping the hybrid recurrent state out of the critical path. FLARE-4B and FLARE-9B are slower in absolute throughput, as expected from their larger backbones, but remain competitive with larger dLLM baselines while providing the capability gains reported above. Full throughput results across three benchmarks and \(C\in\{1,4,8\}\) are reported in Table 10 in Appendix 10.

Together, the three subsections above answer the three questions posed at the start of this section: (i) the converted FLARE checkpoints at the 2B/4B/9B scales preserve the capability of their AR sources and exceed substantially larger dLLM baselines under a more constrained conversion budget; (ii) two-stream supervision reaches MFU comparable to, and slightly above, pure-AR training under our hardware-aware kernel implementation; and (iii) a single checkpoint supports both AR-Trust and Diffusion-Trust decoding under one unified SGLang stack, with substantial end-to-end throughput advantages over other dLLMs at high concurrency.

6 Conclusion, Limitations, and Future Work↩︎

6.0.0.1 Conclusion.

We presented FLARE, a systematic recipe for converting strong hybrid-attention AR models into capable, high-throughput diffusion LLMs. The work contributes several mutually reinforcing designs: a token-balanced clean/noisy two-stream training objective that unifies AR next-token supervision and block-diffusion denoising in one forward pass; a document-packed clean/noisy attention-mask design that supplies both causal and block-bidirectional training signal without cross-document leakage; a study of transfer onto hybrid-attention backbones, where we work out how to realize this objective on mixed softmax/linear-attention architectures through recurrent-state scheduling and matching training kernels, reaching training efficiency comparable to pure-AR training; and a unified inference system in which a single checkpoint supports multiple diffusion decoding modes, from parallel denoising to causal verified decoding. On top of these, our controlled study further identifies transfer-data quality as one of the key factors for capability preservation. Together, these results show that, under a modest conversion budget of roughly 10B tokens, AR-to-dLLM transfer can both retain the source model’s capability and deliver real parallel-decoding throughput, indicating that practical dLLMs are limited less by the decoding algorithm itself than by the joint design of objective, data, and inference system.

6.0.0.2 Limitations.

We group the limitations of this work into three points. (i) Training overhead. Block-diffusion training concatenates a clean and a noisy view of every sequence into a single \(2L\)-length input under a custom attention mask, which roughly doubles the per-step compute and memory relative to size-matched AR training and makes long-context training more expensive; the cost is amplified on hybrid backbones, where linear-attention layers require the noisy stream to read clean intermediate recurrent states at every block boundary — the root cause of the residual training-efficiency gap in Section 5.2. Our dedicated kernels mitigate but do not eliminate this underlying \(2\times\) overhead. (ii) Residual gap to the source model. Even with our best data mix, FLARE still trails its AR source model on several benchmarks. This is consistent with source-distribution shift: continuing SFT on long-CoT traces from external teachers (DeepSeek-R1, Qwen3, GPT-OSS-120B) moves the output distribution away from the post-trained Qwen3.5 seed. Notably, more aggressive filtering via our automatic IFD-based selection pipeline (Appendix 9.5) does not close this gap, indicating that the bottleneck is a distribution mismatch between the transfer data and the source model rather than insufficient data filtering, and pointing to data more tightly aligned with both the source AR teacher and the diffusion objective. (iii) Limited scale and post-training scope. We validate FLARE only on dense checkpoints below 10B parameters and a single SFT stage; MoE backbones and post-training beyond SFT (e.g., reinforcement learning) remain untested, and whether they preserve capability under the same low conversion budget is open.

6.0.0.3 Future work.

These limitations outline open directions for the next generation of dLLM systems [90]. The most fundamental is to move beyond the concatenated two-stream formulation toward single-stream training in which clean and noisy supervision share one forward pass, lowering this overhead along two complementary routes: the backbone side, as in DiffuMamba [91], and the training-cost side, as in Orthrus [92] (tuning only \(\sim\)​16% of parameters). On the data side, the next step is transfer data more tightly aligned with the source-model distribution and the diffusion objective — e.g., diffusion-friendly traces harvested from the same-family AR teacher and matched to the block structure — rather than only finer filtering of existing external corpora. On the scaling side, the recipe can be used to up-cycle dense AR checkpoints into MoE-dLLMs, where LLaDA-MoE [93] already shows feasibility, raising the question of how expert routing interacts with the diffusion mask. Finally, closing the residual quality gap will likely require RL: the per-sequence likelihood of a masked-diffusion model is intractable, so policy-gradient methods need ELBO surrogates [94]–[96], and FLARE’s multiple decoding paths raise the further question of which path to roll out under and how to credit-assign reward across the clean and noisy streams.

7 Hardware-Aware Kernels for Hybrid-Backbone Diffusion Training↩︎

This appendix specifies the kernel-level implementation of the two-stream recurrence defined in Section 3.2: its forward computation on the hybrid GDN backbone, its gradients, and the shipped implementation. A GDN layer contains two sub-components that both require re-implementation under FLARE’s two-stream training: the chunkwise-parallel GDR recurrence, and the width-\(W\) depthwise 1D Causal ShortConv that feeds it. The rest of the appendix proceeds as follows.

• §7.1 states the three challenges that separate FLARE’s two-stream training from single-stream training.

• §7.2 develops the GDR kernels: it reviews the single-stream chunkwise form, lifts it to two streams via the block-boundary clean seed and block-level chain rule, derives the backward primitives, and specifies and compares the two implementation routes.

• §7.3 maps the same construction onto the 1D Causal ShortConv at lower cost, again as two implementation routes.

• §7.4 adds the document-level guards that keep packed samples isolated under recurrent-state propagation.

• §7.5 reports the end-to-end training-MFU gains of the shipped kernel stack on FLARE-2B.

7.1 Why the two-stream case needs a new kernel↩︎

Three challenges that FLARE’s training setup imposes on both sub-components of a GDN layer (the GDR recurrence and the 1D Causal ShortConv) force the design choices of §7.2–§7.4. Standard chunkwise-parallel GDR kernels and standard causal 1D convolution kernels, as shipped by single-stream linear-attention implementations, provide no mechanism for these requirements, since each assumes a single token-causal stream.

7.1.0.1 Challenge 1: clean and noisy streams need asymmetric visibility.

FLARE’s objective (Eq. 4 ) couples two streams: clean tokens are token-causal over the whole clean stream, whereas a noisy token at position \(\ell\) in block \(b\) sees only (i) noisy tokens in the same block and (ii) clean tokens strictly before the block start \((b{-}1)B\). Both sub-components must realize this asymmetry. The GDR seeds the noisy recurrence on block \(b\) from the block-boundary clean state \(\mathbf{S}_{(b-1)B}\); since this state lies inside a chunk whenever the block size \(B\) differs from the chunk size \(C\), a correct implementation must compute \(L/B\) mid-chunk clean states that the single-stream kernel never produces. The 1D Causal ShortConv instead splits its receptive field: a noisy output at offset \(j \in \{0, \ldots, B{-}1\}\) from the block start reads its first \(j{+}1\) lags from the noisy stream and the remaining \(W{-}1{-}j\) lags from the clean stream, whereas a standard causal convolution reads all \(W\) lags from one stream.

7.1.0.2 Challenge 2: document packing requires state resets.

For throughput, multiple documents are packed into each \(L\)-token sequence, with per-document starts encoded as \(\mathbf{cu} = (c_0, c_1, \ldots, c_{N_{\text{doc}}})\); the packed run must reproduce the per-document single-sequence result. This forces a document-level mask to compose with both the two-stream mask and the chunkwise state each sub-component tracks. For the GDR, the noisy recurrence’s initial state at each document’s first block must be zeroed rather than inherited from the previous document’s chunk-boundary state, and the backward’s cross-chunk scan must zero its state-gradient hand-off at document boundaries. For the 1D Causal ShortConv, any lag read crossing a document boundary must be masked to zero instead of reading the previous document’s trailing tokens. Single-stream kernels need none of these guards, since a plain token-causal mask packed with a document-causal mask already forbids cross-document dependencies.

7.1.0.3 Challenge 3: composing three granularities in the backward.

Realizing the Challenge 1 visibility is cheap in the forward— either \(L/B\) extra bytes in HBM or \(\mathcal{O}(B)\) register-replay steps per block for the GDR, and negligible for ShortConv (whose “state” is a raw \(W{-}1\)-token slice). The difficulty concentrates in the backward. A standard chunk-parallel clean backward carries a single cross-chunk gradient \(\mathbf{dh}_{[c]}\equiv \partial\mathcal{L}/\partial\mathbf{S}_{[c]}\) through a matmul scan, delivering all per-token clean gradients at tensor-core throughput. In the two-stream case the noisy backward emits two further objects the chunk-level scan cannot directly accept: per-block clean-transition gradients (§7.2.3.4) and per-block initial-state gradients \(\mathbf{d}\tilde{\mathbf{S}}_{b,\text{init}}\). Routing them into the scan without breaking its parallelism requires composing three granularities of reverse recurrence—token within a block, block within a chunk, and chunk within the sequence—in the correct order (§7.2.3). The ShortConv backward has only two of these granularities, because its noisy-side “state” is a static clean-stream slice rather than a recurrence terminus, which makes it cheaper than the GDR ’s (§7.3.2). The routes of §7.2.5–§7.2.6 are two schedules of this composition onto a GPU; the resulting block-by-block backward dataflow is detailed in Figure 14.

7.1.0.4 AR-only training does not trigger these challenges.

AR-only training of a GDN layer runs a single stream with uniform token-causal visibility. Challenge 1 does not arise because there is a single stream and no block-level granularity. Challenge 2 reduces to standard document-causal attention masking, which composes with the token-causal per-stream mask without any additional kernel-side guard. Challenge 3 reduces to a single cross-chunk state-gradient tensor at chunk-level granularity, which the standard chunkwise-parallel kernels of [20], [75] already handle. FLARE’s block-seeded noisy-stream objective is what introduces the machinery described below.

7.1.0.5 Notation.

Table 3 collects the symbols used in the remainder of this appendix. Positions in a stream are indexed by \(t\) (clean) or \(\ell\) (noisy). The lengths \(B\), \(C\), \(L\) and derived counts \(M\), \(N_C\) satisfy \(L = C \cdot N_C\) and \(C = B \cdot M\); we use \(C{=}64\) throughout.

7.2 Gated Delta Rule: two-stream training and implementation routes↩︎

This subsection develops the chunkwise-parallel training routes for the GDR, the recurrence sub-component of the GDN layer. The objects and algorithms introduced here (chunk-boundary states, WY-corrected values, block-level chain rule, Route I and Route II) constitute the main technical content of the appendix; §7.3 reuses the same structure for the 1D Causal ShortConv at lower cost.

7.2.1 Single-stream chunkwise training↩︎

Both routes reuse the standard chunkwise-parallel training form of a single-stream GDR on the clean side unchanged. We restate it here to fix names for the four objects the clean-side forward produces and both routes reuse: the clean output, the chunk-boundary states, the chunk-corrected (WY) values, and the cumulative log-gates.

7.2.1.1 Per-step recurrence.

Expanding the GDR of Eq. 3 into its elementary per-step state and output updates yields the form used in both optimization and differentiation: \[\label{eq:per-step-fwd} \begin{align} & \mathbf{S}^{\text{pre}}_t = \exp\,(g_t)\,\mathbf{S}_{t-1}, & \text{(gating decay)}\\ & u_t = \mathbf{v}_t - (\mathbf{S}^{\text{pre}}_t)^{\!\top}\mathbf{k}_t, & \text{(prediction error)}\\ & \boldsymbol{\delta}_t = \beta_t\,u_t, & \text{(delta update)}\\ & \mathbf{S}_t = \mathbf{S}^{\text{pre}}_t + \mathbf{k}_t\boldsymbol{\delta}_t^{\!\top}, & \text{(outer-product write)}\\ & \mathbf{o}_t = s\cdot\mathbf{S}_t^{\!\top}\mathbf{q}_t, & \text{(output)} \end{align}\tag{13}\] where \(\mathbf{S}^{\text{pre}}_t \in \mathbb{R}^{d_k \times d_v}\) is the state after the gating decay but before the delta update, \(u_t \in \mathbb{R}^{d_v}\) is the per-step prediction error between the target value \(\mathbf{v}_t\) and the value currently associated with \(\mathbf{k}_t\) under \(\mathbf{S}^{\text{pre}}_t\), \(\boldsymbol{\delta}_t = \beta_t u_t \in \mathbb{R}^{d_v}\) is the \(\beta\)-weighted correction written into the state by the outer-product update, \(\mathbf{S}_t \in \mathbb{R}^{d_k \times d_v}\) is the state after the update, \(\mathbf{o}_t \in \mathbb{R}^{d_v}\) is the per-token output, and \(s = 1/\sqrt{d_k}\) is the standard query-key scaling. The inputs \((g_t, \beta_t, \mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t)\) are produced by the linear projections and ShortConv branches of the GDN layer and are the same as in Eq. 3 . §7.2.3.1 differentiates the five lines.

7.2.1.2 Chunkwise-parallel form.

A token-by-token implementation of Eq. 13 is serial in the sequence dimension and underutilizes tensor cores. A fully materialized \(L \times L\) form reaches tensor cores but incurs \(\mathcal{O}(L^2)\) cost and loses the linear-attention advantage. The chunkwise-parallel form recovers both: split the sequence into \(N_C = L/C\) chunks of size \(C\) (we set \(C = 64\) to match tensor-core tile sizes), process each chunk with dense matrix multiplication, and carry a single compact state \(\mathbf{S}_{[c]} \in \mathbb{R}^{d_k \times d_v}\) across chunks. Two pieces remain: how each chunk is processed in parallel despite the delta rule’s within-chunk dependence, and how the compact state is advanced across chunks; the next two paragraphs address them in turn.

7.2.1.3 Chunk-corrected values.

The delta-rule correction at each token depends on the partial state that earlier tokens in the same chunk have already written, so running Eq. 13 token-by-token within a chunk incurs \(C\)-step serial depth. The WY representation of the delta rule [20], [64] removes this depth by replacing each raw value \(\mathbf{v}_t\) with a chunk-corrected (or effective) value \(\hat{\mathbf{v}}_t\), computed from \(\{(\mathbf{k}_t, \mathbf{v}_t, \beta_t, g_t)\}_{t\in\text{chunk}(c)}\) via a small triangular solve, so that the chunk’s contribution to the state reduces to a single outer-product sum. We use \(\hat{\mathbf{v}}_t\) throughout; in particular, the chunk-level state recurrence below is written directly in terms of \(\hat{\mathbf{v}}_t\).

7.2.1.4 Chunk-level state recurrence.

Let \(G_t = \sum_{i \le t} g_i\) denote the per-token cumulative log-gate on the clean stream, and let \(\mathbf{S}_{[c]} \equiv \mathbf{S}_{cC}\) denote the clean state at the boundary between chunks \(c-1\) and \(c\). In terms of \(\hat{\mathbf{v}}\) and \(G\), the single-stream GDR state recurrence over one chunk \(c\) admits the closed form \[\label{eq:chunk-state} \mathbf{S}_{[c+1]} \;=\; \exp\,\bigl(G_{(c+1)C}-G_{cC}\bigr)\,\mathbf{S}_{[c]} \;+\; \sum_{t\in\text{chunk}(c)} \exp\,\bigl(G_{(c+1)C}-G_t\bigr)\,\mathbf{k}_t\,\hat{\mathbf{v}}_t^{\!\top},\tag{14}\] where \(\exp\,(G_{(c+1)C}-G_{cC})\) is the scalar decay factor produced by compounding the \(C\) per-token gates of chunk \(c\), and \(\exp\,(G_{(c+1)C}-G_t)\) is the residual decay from position \(t\) to the chunk end. Within-chunk work reduces to a dense sum of outer products (tensor-core friendly), and the only cross-chunk serial dependency is the single \(\mathbf{S}_{[c]} \to \mathbf{S}_{[c+1]}\) hand-off. The chunk-level output \(\mathbf{O}_{[c]}\) admits a similar closed form from the chunk’s queries, \(\mathbf{S}_{[c]}\), and in-chunk keys; since neither route modifies that kernel we omit it, though §7.2.3.3 differentiates its state-side contribution. Stacking \(\mathbf{O}_{[c]}\) over all chunks gives the clean-stream output \(\mathbf{o}\) listed below.

7.2.1.5 Clean-side forward outputs.

The clean-side forward produces four objects, reused unchanged by both routes:

1. \(\mathbf{o}\), the clean-stream output (the standard chunkwise attention output);

2. \(\{\mathbf{S}_{[c]}\}_{c=0}^{N_C}\), chunk-boundary states from the scan in Eq. 14 ;

3. \(\{\hat{\mathbf{v}}_t\}_{t=1}^{L}\), chunk-corrected clean values from the WY step;

4. \(\{G_t\}_{t=1}^{L}\), per-position cumulative log-gates.

The route-specific work begins after these outputs: it concerns only the construction of the block-boundary clean states \(\tilde{\mathbf{S}}_{b,\text{init}}\) that seed the noisy stream, which §7.2.2 takes up next.

7.2.2 Two-stream recurrence and block-level chain rule↩︎

The clean side of §7.2.1 is inherited unchanged from standard single-stream GDR training; the FLARE-specific work supplies the noisy stream with its block-boundary clean seed. This subsection defines the seed, specifies how the noisy recurrence consumes it, and derives the block-level chain rule that both routes implement.

7.2.2.1 Clean-transition map over one block.

Specialising Eq. 14 to a single block of \(B\) tokens defines the clean-transition map. For block \(b\) spanning positions \((b-1)B+1, \ldots, bB\): \[\label{eq:f-blk} f^{\text{blk}}_b(\mathbf{S}) \;=\; \exp\,\bigl(G_{bB}-G_{(b-1)B}\bigr)\,\mathbf{S} \;+\; \boldsymbol{\Delta}_b, \qquad \boldsymbol{\Delta}_b \;=\!\!\!\sum_{t=(b-1)B+1}^{bB}\!\!\! \exp\,\bigl(G_{bB}-G_t\bigr)\,\mathbf{k}_t\,\hat{\mathbf{v}}_t^{\!\top},\tag{15}\] where \(f^{\text{blk}}_b : \mathbb{R}^{d_k \times d_v} \to \mathbb{R}^{d_k \times d_v}\) is the block-\(b\) clean-transition map that advances the clean state by \(B\) tokens, \(\exp\,(G_{bB}-G_{(b-1)B})\) is the block’s cumulative gate, and \(\boldsymbol{\Delta}_b \in \mathbb{R}^{d_k \times d_v}\) is the block’s outer-product write in terms of raw clean keys and chunk-corrected clean values.

7.2.2.2 Block-boundary clean states.

Inside chunk \(c\), the \(M\) block-boundary clean states \(\mathbf{S}_{cM \cdot B + B}, \mathbf{S}_{cM \cdot B + 2B}, \ldots, \mathbf{S}_{(c+1)C}\) are obtained by composing \(f^{\text{blk}}\) starting from the chunk-boundary state \(\mathbf{S}_{[c]}\); applying \(f^{\text{blk}}\) all \(M\) times recovers \(\mathbf{S}_{[c+1]}\) and Eq. 14 at the next chunk boundary. The two routes agree on the definition of these \(L/B\) states and differ only in the stage at which each state is materialized.

7.2.2.3 Noisy block forward.

For block \(b\), the noisy recurrence starts from the corresponding clean block-end state as its initial state, then runs Eq. 13 on the noisy tokens of that block only: \[\label{eq:noisy-block-fwd} \tilde{\mathbf{S}}_{b,\text{init}} \;\mathrel{:=}\; \mathbf{S}_{(b-1)B}, \qquad \tilde{\mathbf{S}}_\ell \;=\; \exp\,(\tilde{g}_\ell)\,\tilde{\mathbf{S}}_{\ell-1} \;+\; \tilde{\mathbf{k}}_\ell\,\tilde{\boldsymbol{\delta}}_\ell^{\!\top},\tag{16}\] for \(\ell = (b-1)B+1, \ldots, bB\), where \(\tilde{\mathbf{S}}_{b,\text{init}} \in \mathbb{R}^{d_k \times d_v}\) is block \(b\)’s noisy-recurrence initial state (seeded from the clean state at the block boundary), \(\tilde{\mathbf{S}}_\ell \in \mathbb{R}^{d_k \times d_v}\) is the noisy state after noisy position \(\ell\), and \((\tilde{g}_\ell, \tilde{\beta}_\ell, \tilde{\mathbf{q}}_\ell, \tilde{\mathbf{k}}_\ell, \tilde{\mathbf{v}}_\ell)\) are the noisy per-token inputs (produced by the same linear projections and ShortConv branches as their clean counterparts, but applied to the noisy stream). The noisy per-token delta update \(\tilde{\boldsymbol{\delta}}_\ell = \tilde{\beta}_\ell\,\tilde{u}_\ell \in \mathbb{R}^{d_v}\) is defined via the noisy prediction error \(\tilde{u}_\ell = \tilde{\mathbf{v}}_\ell - (\tilde{\mathbf{S}}^{\text{pre}}_\ell)^{\!\top}\tilde{\mathbf{k}}_\ell\) with \(\tilde{\mathbf{S}}^{\text{pre}}_\ell = \exp\,(\tilde{g}_\ell)\,\tilde{\mathbf{S}}_{\ell-1}\), reusing the five-line expansion of Eq. 13 on the noisy stream without restating it. The noisy output at every position in block \(b\) is read from the block-end state rather than the running state: \(\tilde{\mathbf{o}}_\ell = s\,\tilde{\mathbf{S}}_{bB}^{\!\top}\tilde{\mathbf{q}}_\ell\) for all \(\ell \in \{(b{-}1)B{+}1, \ldots, bB\}\). This single-shared readout state reproduces the bidirectional noisy-to-noisy region of the training mask (Eq. 6 ). Between blocks the noisy state is reset; within a block it advances exactly \(B\) steps. Noisy blocks are mutually independent because each is seeded from its own clean state; all block-to-block information flow passes through the clean stream.

7.2.2.4 Block-level chain rule.

Differentiating Eq. 15 yields the chain rule that both routes reuse. For a block \(b\) inside chunk \(c\), let \[\mathbf{dh}_b \;\equiv\; \partial\mathcal{L}/\partial\mathbf{S}_{bB}\] denote the clean state gradient at the end of block \(b\), accumulated so far from later blocks in the same chunk, and let \(\mathbf{d}\tilde{\mathbf{S}}_{b,\text{init}}\) denote the gradient produced by reverse-sweeping the noisy recurrence of Eq. 16 on block \(b\) alone. The block-level recursion is then \[\label{eq:dh-block-recursion} \mathbf{dh}_{b-1} \;=\; \mathbf{d}\tilde{\mathbf{S}}_{b,\text{init}} \;+\; \exp\,\bigl(G_{bB}-G_{(b-1)B}\bigr)\, \mathbf{dh}_b.\tag{17}\] The first term is block \(b\)’s direct contribution to the gradient at its seed (the clean state at the end of block \(b{-}1\)). The second term is the gate-decayed gradient from blocks \(b{+}1, b{+}2, \ldots\) propagated through \(f^{\text{blk}}_b\)’s scalar decay. Isolating the within-chunk part of this scan by seeding the recursion with \(\mathbf{dh}_{cM{+}M} = \mathbf{0}\) (dropping external contributions to chunk \(c\)) and iterating \(M\) times yields \(\mathbf{dh}_{cM}\), the within-chunk contribution \(\mathbf{dh}^{\text{inject}}_{[c]}\). The chunk-level scan of Eq. 20 subsequently adds the two external contributions (from later chunks and from chunk \(c\)’s clean-output gradient) on top of this injection.

7.2.3 Backward primitives↩︎

Three primitives appear repeatedly in both routes: a per-token backward inside one block (§7.2.3.1); a chunk-level backward across chunks on the clean side (§7.2.3.3); and a clean-transition backward that differentiates \(f^{\text{blk}}_b\) (§7.2.3.4). We derive each once and invoke them by name in Routes I and II.

7.2.3.1 Per-step backward.

The per-step backward reverses a single token of Eq. 13 ; each reverse sweep over a noisy block consists of \(B\) consecutive invocations of this primitive. Let \(\mathbf{dS}_t = \partial\mathcal{L}/\partial\mathbf{S}_t\), \(\mathbf{do}_t = \partial\mathcal{L}/\partial\mathbf{o}_t\), \(\mathbf{dS}^{\text{pre}}_t = \partial\mathcal{L}/\partial\mathbf{S}^{\text{pre}}_t\), and \(\mathbf{d}u_t = \partial\mathcal{L}/\partial u_t\). Because \(\mathbf{S}_t\) appears both in \(\mathbf{o}_t\) and, through the subsequent token, in \(\mathbf{S}_{t+1}\), its gradient is an accumulator: by the time \(\mathbf{dS}_{t-1}\) is computed, \(\mathbf{dS}_t\) has already received the contribution from \(\mathbf{S}_{t+1}\) handed back by the reverse step on token \(t{+}1\). The equations below therefore use ‘\(\mathrel{+}=\)’ (and in Eqs. 21 –22 also ‘\(\mathrel{-}=\)’) in the programming sense: \(x\mathrel{+}=y\) denotes \(x \leftarrow x + y\). Differentiating the five lines of Eq. 13 in reverse yields \[\label{eq:per-step-bwd} \begin{align} & \mathbf{dq}_t = s\,\mathbf{S}_t\mathbf{do}_t, & \text{(from \mathbf{o}_t)}\\ & \mathbf{dS}_t \mathrel{+}= s\,\mathbf{q}_t\,\mathbf{do}_t^{\!\top}, & \text{(from \mathbf{o}_t; add into \mathbf{dS}_t)}\\ & \mathbf{d}\boldsymbol{\delta}_t = \mathbf{dS}_t^{\!\top}\mathbf{k}_t, \;\;\mathbf{dk}_t^{\text{write}} = \mathbf{dS}_t\,\boldsymbol{\delta}_t, & \text{(from the outer-product write)}\\ & \mathbf{d}u_t = \beta_t\,\mathbf{d}\boldsymbol{\delta}_t, \;\;\mathbf{dv}_t = \mathbf{d}u_t, \;\;d\beta_t = \mathbf{d}\boldsymbol{\delta}_t^{\!\top}u_t, & \text{(from the delta update)}\\ & \mathbf{dk}_t^{\text{pred}} = -\mathbf{S}^{\text{pre}}_t\,\mathbf{d}u_t, \;\;\mathbf{dk}_t = \mathbf{dk}_t^{\text{write}} + \mathbf{dk}_t^{\text{pred}}, & \text{(from the prediction error)}\\ & \mathbf{dS}^{\text{pre}}_t = \mathbf{dS}_t - \mathbf{k}_t\,\mathbf{d}u_t^{\!\top}, & \text{(splits \mathbf{S}_t = \mathbf{S}^{\text{pre}}_t + \mathbf{k}_t\boldsymbol{\delta}_t^{\!\top})}\\ & dg_t = \bigl\langle \mathbf{dS}^{\text{pre}}_t,\;\mathbf{S}^{\text{pre}}_t\bigr\rangle_{\mathrm{F}}, & \text{(from \mathbf{S}^{\text{pre}}_t = \exp\,(g_t)\,\mathbf{S}_{t-1})}\\ & \mathbf{dS}_{t-1} = \exp\,(g_t)\,\mathbf{dS}^{\text{pre}}_t, & \text{(gating reverse recurrence)} \end{align}\tag{18}\] where \(\mathbf{d}\boldsymbol{\delta}_t = \partial\mathcal{L}/\partial\boldsymbol{\delta}_t \in \mathbb{R}^{d_v}\) is the gradient with respect to the delta update, \(\mathbf{dk}_t^{\text{write}}, \mathbf{dk}_t^{\text{pred}} \in \mathbb{R}^{d_k}\) are the two contributions to \(\mathbf{dk}_t\) from the outer-product write and the prediction error, respectively, and \(\langle\cdot,\cdot\rangle_{\mathrm{F}}\) denotes the Frobenius inner product. The key gradient sums the two contributions because \(\mathbf{k}_t\) enters both the outer-product write and the prediction error. Denote this primitive by PerStepBwd: it accepts \((\mathbf{dS}_t, \mathbf{do}_t, \mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t, g_t, \beta_t, \mathbf{S}^{\text{pre}}_t, \mathbf{S}_t)\) and returns \((\mathbf{dq}_t, \mathbf{dk}_t, \mathbf{dv}_t, dg_t, d\beta_t, \mathbf{dS}_{t-1})\). Eq. 18 assumes the token-causal readout \(\mathbf{o}_t = s\,\mathbf{S}_t^{\!\top}\mathbf{q}_t\): each per-step state feeds exactly one output, so the output-side gradient enters the state accumulator at the same step. This is correct for the clean stream; the noisy stream uses a different output rule, which we cover next.

7.2.3.2 Noisy-stream backward under block-end readout.

The noisy-block forward of Eq. 16 reads every noisy output in block \(b\) from the same end-of-block state \(\tilde{\mathbf{S}}_{bB}\), so the mapping from outputs to states is many-to-one instead of one-to-one. Differentiating this readout gives a block-level output-side contribution that lives entirely on \(\tilde{\mathbf{S}}_{bB}\)’s gradient, plus a per-step query gradient that uses the same single state: \[\label{eq:noisy-readout-bwd} \mathbf{d}\tilde{\mathbf{S}}_{bB} \;=\; s\!\!\sum_{\ell=(b-1)B+1}^{bB}\!\!\tilde{\mathbf{q}}_\ell\,\mathbf{d}\tilde{\mathbf{o}}_\ell^{\!\top}, \qquad \mathbf{d}\tilde{\mathbf{q}}_\ell \;=\; s\,\tilde{\mathbf{S}}_{bB}\,\mathbf{d}\tilde{\mathbf{o}}_\ell \quad\text{for all } \ell\in\{(b-1)B+1,\ldots,bB\}.\tag{19}\] A reverse sweep over the noisy block then runs \(B\) consecutive PerStepBwd calls with \(\mathbf{d}\tilde{\mathbf{o}}_\ell\) set to \(\mathbf{0}\): the per-step output terms (the \(\mathbf{dq}_t\) and \(\mathbf{dS}_t \mathrel{+}{=} s\,\mathbf{q}_t\,\mathbf{do}_t^{\!\top}\) lines of Eq. 18 ) are already supplied by Eq. 19 , and the remaining lines (gating, outer-product write, prediction error, and gating reverse recurrence) are unchanged. Throughout this appendix, “reverse sweep over a noisy block of \(B\) tokens” refers to this procedure: initialize \(\mathbf{d}\tilde{\mathbf{S}}_{bB}\) with Eq. 19 , emit \(\mathbf{d}\tilde{\mathbf{q}}_\ell\) directly from Eq. 19 , then invoke PerStepBwd with \(\mathbf{do} = \mathbf{0}\) for each of the \(B\) steps in reverse.

7.2.3.3 Chunk-level clean backward.

The chunk-level backward differentiates Eq. 14 together with the clean-output kernel that emits \(\mathbf{O}_{[c]}\). Its only cross-chunk serial step is a reverse matmul scan over chunk-boundary state gradients \(\mathbf{dh}_{[c]} \equiv \partial\mathcal{L}/\partial\mathbf{S}_{[c]}\): \[\label{eq:chunk-bwd-scan} \mathbf{dh}_{[c]} \;=\; \underbrace{\mathbf{dh}^{\text{inject}}_{[c]}}_{\text{external hand-off}} \;+\; \underbrace{\mathcal{C}_{[c]}}_{\text{in-chunk clean-output contribution}} \;+\; \exp\,\bigl(G_{(c+1)C}-G_{cC}\bigr)\,\mathbf{dh}_{[c+1]},\tag{20}\] where \(\mathbf{dh}^{\text{inject}}_{[c]} \in \mathbb{R}^{d_k \times d_v}\) is a per-chunk hand-off slot that accepts external state-gradients and adds them into the scan (written by Route I’s fill-in backward, §7.2.5, or by Route II’s fused kernel, §7.2.6), and \(\mathcal{C}_{[c]} \in \mathbb{R}^{d_k \times d_v}\) is the in-chunk clean-output contribution obtained by differentiating \(\mathbf{O}_{[c]}\) through \(\mathbf{S}_{[c]}\) via a dense matmul over the \(C\) tokens of chunk \(c\). \(\mathcal{C}_{[c]}\) is identical to the AR baseline’s clean-side computation and is not modified by either route; the two-stream-specific content is routed entirely through \(\mathbf{dh}^{\text{inject}}_{[c]}\). Denote this primitive by ChunkBwd.

7.2.3.4 Clean-transition backward.

The clean-transition backward links the noisy-side reverse sweep to the clean-side cross-chunk scan: it converts the state gradient \(\mathbf{dh}_b\) at the end of block \(b\) into per-token clean gradients and a propagated block-initial state gradient. It is the block-level analog of §7.2.3.3. We differentiate \(f^{\text{blk}}_b\) (Eq. 15 ) with respect to its clean-side inputs \(\bigl(\mathbf{k}_t, \hat{\mathbf{v}}_t, G_t\bigr)_{t\in\text{block}(b)}\) and the incoming state \(\mathbf{S}\). Let \(\rho_t \equiv \exp\,(G_{bB}-G_t)\) denote the per-token decay factor inside block \(b\), and let \(\mathbf{dh}_b \equiv \partial\mathcal{L}/\partial\mathbf{S}_{bB}\) denote the incoming state gradient.

7.2.3.5 Gradient through the scalar decay.

\(f^{\text{blk}}_b\)’s leading term \(\exp\,(G_{bB}-G_{(b-1)B})\,\mathbf{S}\) propagates \(\mathbf{dh}_b\) to \(\mathbf{S}\) with factor \(\exp\,(G_{bB}-G_{(b-1)B})\); this factor is exactly the multiplicative factor in the block-level recursion of Eq. 17 , and is how the recursion is derived. The same decay also yields gradients into \(G_{bB}\) and \(G_{(b-1)B}\): \[\label{eq:dalpha} \begin{align} & dG_{bB} \mathrel{+}= \exp\,\bigl(G_{bB}-G_{(b-1)B}\bigr)\,\bigl\langle \mathbf{dh}_b,\;\mathbf{S}\bigr\rangle_{\mathrm{F}}, & \text{(contribution to dG_{bB})}\\ & dG_{(b-1)B} \mathrel{-}= \exp\,\bigl(G_{bB}-G_{(b-1)B}\bigr)\,\bigl\langle \mathbf{dh}_b,\;\mathbf{S}\bigr\rangle_{\mathrm{F}}. & \text{(contribution to dG_{(b-1)B})} \end{align}\tag{21}\] The two writes have equal magnitude and opposite signs because \(\partial(G_j - G_i)/\partial G_j = 1\) and \(\partial(G_j - G_i)/\partial G_i = -1\). Since each \(dG\) accumulator receives contributions from several sources (this block’s scalar-decay term of Eq. 21 , this block’s \(\boldsymbol{\Delta}\) term below, and any later calls that hit the same \(G\)), we keep the ‘\(\mathrel{+}=\)’ / ‘\(\mathrel{-}=\)’ form throughout §7.2.3.4 to make the accumulator semantics explicit.

7.2.3.6 Gradient through the outer-product write.

The write term \(\boldsymbol{\Delta}_b = \sum_t \rho_t\,\mathbf{k}_t\,\hat{\mathbf{v}}_t^{\!\top}\) contributes, for each \(t\) in block \(b\): \[\label{eq:dDelta} \begin{align} & \mathbf{dk}_t = \mathbf{dh}_b\bigl(\hat{\mathbf{v}}_t\,\rho_t\bigr), & \text{(clean key)}\\ & \mathbf{d}\hat{\mathbf{v}}_t = \rho_t\,\mathbf{dh}_b^{\!\top}\mathbf{k}_t, & \text{(chunk-corrected value)}\\ & \sigma_t := \rho_t\,\mathbf{k}_t^{\!\top}\mathbf{dh}_b\,\hat{\mathbf{v}}_t, \;\;dG_{bB}\mathrel{+}= \sigma_t, \;\;dG_t\mathrel{-}= \sigma_t, & \text{(cumulative gates)} \end{align}\tag{22}\] where \(\mathbf{dk}_t \in \mathbb{R}^{d_k}\) and \(\mathbf{d}\hat{\mathbf{v}}_t \in \mathbb{R}^{d_v}\) are the per-token gradients of the block’s outer-product write term with respect to the clean key and chunk-corrected clean value, and \(\sigma_t \in \mathbb{R}\) is the scalar obtained by contracting \(\mathbf{dh}_b\) with the rank-one update \(\mathbf{k}_t\hat{\mathbf{v}}_t^{\!\top}\) and multiplying by \(\rho_t\); it is the common magnitude of the two \(dG\) writes. The two \(G\)-writes have opposite signs because \(G_{bB}\) and \(G_t\) enter \(\rho_t\) with opposite signs. Summed over \(t \in \text{block}(b)\) (and over all blocks and all contribution sources), these give the final per-token clean-side gradients \((\mathbf{dk}_t, \mathbf{d}\hat{\mathbf{v}}_t, dG_t)\) of block \(b\).

7.2.3.7 Summary.

Denote the combined sub-procedure by CleanTransBwd: it accepts \((\mathbf{dh}_b, \mathbf{S}, \{(\mathbf{k}_t, \hat{\mathbf{v}}_t, G_t)\}_{t\in\text{block}(b)})\) and returns \((\mathbf{dk}_t, \mathbf{d}\hat{\mathbf{v}}_t, dG_t)\) for every token in block \(b\). Its outputs are partial-gradient slots: the downstream chunk-level clean backward (ChunkBwd, together with the WY backward that converts \(\mathbf{d}\hat{\mathbf{v}}\) into \((\mathbf{dv}, \mathbf{dk}, d\beta)\)) completes backpropagation into the final clean parameter gradients.

7.2.3.8 Gradient map.

Figure 14 traces the composition of the three primitives on a single noisy token. Both routes produce the same tree and differ only in which kernel computes which subtree.

7.2.4 Overview of the two implementation routes↩︎

Both routes compute identical gradients and differ only in (i) how the block-boundary clean states \(\mathbf{S}_{(b-1)B}\) are delivered on the forward and (ii) how block-level gradients are routed into the cross-chunk scan. Figure 15 contrasts their dataflow; §7.2.5 and §7.2.6 specify the algorithms and §7.2.7 quantifies the trade-off.

The block-boundary clean states \(\mathbf{S}_{(b-1)B}\) follow §7.2.2; both routes ultimately route their gradients into the chunk-boundary slot consumed by ChunkBwd.

7.2.5 Route I: Chunk-then-Refine↩︎

Route I materializes every \(\mathbf{S}_{(b-1)B}\) in HBM prior to the noisy forward. The noisy stream then reduces to \(L/B\) independent block-local recurrences that run in parallel.

1. Clean chunkwise forward. Run §7.2.1 on the clean stream; emit \(\mathbf{o}\), \(\{\mathbf{S}_{[c]}\}\), \(\{\hat{\mathbf{v}}_t\}\), and \(\{G_t\}\).

2. Block-state fill-in (forward). For each chunk \(c\), apply \(f^{\text{blk}}_{cM+1}, \ldots, f^{\text{blk}}_{cM+M-1}\) (Eq. 15 ) starting from \(\mathbf{S}_{[c]}\), and write every \(\mathbf{S}_{(b-1)B}\) for \(b\) inside chunk \(c\) to HBM. The pass has no tensor-core work; it is a per-block recurrent scan whose only purpose is to populate the tensor \(\{\mathbf{S}_{(b-1)B}\}_{b=1}^{L/B}\).

3. Block-local noisy forward. Treat the noisy stream as \(L/B\) independent length-\(B\) sequences; each starts from its fill-in state \(\mathbf{S}_{(b-1)B}\), runs Eq. 16 for \(B\) steps, and emits \(\tilde{\mathbf{o}}\) for those positions.

7.2.5.1 Backward.

The backward mirrors the forward as three kernels connected by autograd through two HBM hand-off tensors (Figure [fig:route-bwd], Route I row). Using the primitives of §7.2.3:

Route I’s two costs are an \(L/B\)-sized tensor of block-boundary clean states in HBM and a three-kernel serial backward chain; Route II (§7.2.6) eliminates both, and §7.2.7 quantifies the difference.

7.2.6 Route II: Fused Two-Stream↩︎

Route II collapses forward stages 2–3 (and the corresponding backward stages) into a single fused kernel per direction. The forward kernel launches one program per chunk, loads \(\mathbf{S}_{[c]}\) into registers, and interleaves noisy-block forward passes with applications of \(f^{\text{blk}}\) that advance the in-register clean state. No block-boundary clean state is written to HBM. The backward launches one program per \((c,\, i_V)\) pair, where \(i_V \in \{1, \ldots, \lceil d_v/b_v \rceil\}\) indexes the value-dimension tile, and sweeps the same structure in reverse, invoking the primitives of §7.2.3 as named sub-procedures.

7.2.6.1 Strided clean-state checkpoints.

The backward requires \(\mathbf{S}_{(b-1)B}\) for every block \(b\) in the chunk, but the forward keeps only \(\mathbf{S}_{[c]}\) resident. A direct remedy is to store every block-end clean state to HBM (\(M\) snapshots per chunk, \(L/B\) snapshots per layer per sample), which reproduces Route I’s fill-in tensor and its \(L/B\) footprint. We instead adopt a strided-checkpoint tensor parameterised by the checkpoint stride \(S \in \{1,2,\ldots,M\}\): within each chunk only every \(S\)-th block-boundary clean state is stored, giving \(N_{\text{ckpt}} = M/S\) snapshots per chunk and \(N_C\cdot M/S = L/(BS)\) snapshots per layer per sample. The remaining \((S{-}1)/S\) fraction of block-end states are not persisted; the backward reconstructs each by loading the nearest prior snapshot and applying at most \(S{-}1\) copies of \(f^{\text{blk}}\) in registers (Fig. 19). \(S\) parameterises a time/memory trade-off: the HBM cost scales as \(1/S\) and the per-block replay work scales as \(S{-}1\). Defaults are chosen by block size (§7.2.6.2); on the Qwen3.5-2B shape at \(B{=}1, S{=}16\) the checkpoint tensor is approximately \(128\) MiB, roughly \(16\times\) smaller than Route I’s approximately \(2\) GiB state tensor.

7.2.6.2 Choosing the stride.

\(S\) is chosen so that Route II’s extra HBM stays roughly constant independent of \(B\): since the per-chunk snapshot count is \(N_{\text{ckpt}} = M/S = C/(BS)\), we scale \(S\) with \(M\) so that \(N_{\text{ckpt}}\) stays small (\(\mathcal{O}(1)\)) per chunk. With \(C{=}64\) fixed, the defaults are \(S{=}16\) at \(B{=}1\) (\(N_{\text{ckpt}}{=}4\)), \(S{=}8\) at \(B{=}2\) (\(N_{\text{ckpt}}{=}4\)), \(S{=}2\) at \(B{=}4\) (\(N_{\text{ckpt}}{=}8\)), and \(S{=}\min(8,M)\) otherwise; across these configurations \(N_{\text{ckpt}}\in\{1,4,8\}\) and the checkpoint tensor occupies between \(32\) MiB and \(256\) MiB on the Qwen3.5-2B shape. The per-block replay cost is at most \(S{-}1\) evaluations of \(f^{\text{blk}}\) in registers and remains register-bound because each \(f^{\text{blk}}\) is a single \(\mathbb{R}^{d_k\times d_v}\) update that fits in one tile.

7.2.6.3 Fused forward kernel.

The forward traverses each chunk block-by-block in order: at block \(j\) the in-register clean state serves as the noisy-recurrence initial state; after the \(B\) noisy-token forward steps, the clean state is advanced by \(f^{\text{blk}}\) and the iteration proceeds. Only noisy outputs are written to HBM; the clean output is produced by the separate chunkwise kernel of §7.2.1.

7.2.6.4 Fused backward kernel.

The backward traverses each chunk block-by-block in reverse; the body is given in Alg. 22 and visualized in Fig. 23. All per-block replay, reverse sweep, and clean-transition work is performed in registers; the kernel writes only the single chunk-boundary tensor \(\mathbf{dh}^{\text{inject}}_{[c]}\) to HBM, which ChunkBwd consumes exactly as in Route I.

7.2.6.5 Parallelism structure.

The backward of Fig. 14 groups into three parallelism tiers once Levels 1 and 2 (within a block) are folded together: tier A performs one block-end output contraction plus \(B\) per-step reverse steps inside a block; tier B sweeps \(M\) such blocks inside a chunk via CleanTransBwd and the block-level recursion; and tier C sweeps \(N_C\) chunks across the sequence via ChunkBwd. Route II performs tiers A and B entirely in registers inside the fused kernel and delegates tier C to ChunkBwd, the same chunk-level matmul scan used by the AR baseline. Since tier B is strictly within-chunk, chunks are independent at the fused-kernel level, which enables Route II’s grid of \(N_C\cdot B_{\text{batch}}\cdot H\cdot\lceil d_v/b_v\rceil\) programs; Route I combines tier B with tier C in the same kernel and is therefore restricted to a grid of \(B_{\text{batch}}\cdot H\cdot\lceil d_v/b_v\rceil\) programs.

7.2.7 Implementation comparison↩︎

The two routes produce identical gradients, so the selection between them is governed by resource trade-offs. Route I has fewer moving parts: three small kernels per direction and no strided-checkpoint bookkeeping. Route II reduces peak memory and increases backward parallelism. In the deployed FLARE configuration (\(B < 16\)) Route II dominates on every axis considered below; at \(B \ge 16\) Route I is faster on latency while Route II retains the peak-memory advantage. The three axes are listed in order of importance at small \(B\).

7.2.7.1 Peak memory.

Route I materializes an \(L/B\)-scaling tensor of block-boundary clean states: \(\{\mathbf{S}_{(b-1)B}\}_{b=1}^{L/B}\) is written in forward stage 2 and kept alive until consumed by the block-local noisy forward (and, in general, by the fill-in backward). Its size is \[\text{HBM}_{\text{Route~I, states}} \;=\; \underbrace{L/B}_{\text{grows as B shrinks}} \cdot\,H\cdot d_k\cdot d_v\cdot\text{(bytes per element)} ,\] which at \(B{=}1\) on the Qwen3.5-2B shape is approximately \(2\) GiB per layer per sample in bf16; across the backbone’s \(18\) GDN layers this exhausts the available HBM at the small block sizes used in FLARE training. Route II does not form this tensor: \(\mathbf{S}_{(b-1)B}\) is recomputed in registers from the nearest strided checkpoint and used immediately. The only extra HBM Route II allocates is the strided-checkpoint tensor of Alg. 20, which scales as \(L/(BS)\) and at the \(B{=}1\) default \(S{=}16\) on the same shape is approximately \(128\) MiB, roughly \(16\times\) smaller.

7.2.7.2 Per-chunk parallelism.

Route I’s backward combines cross-chunk propagation with the per-block reverse sweep in a single kernel, so its noisy-backward grid is \(B_{\text{batch}} \cdot H \cdot \lceil d_v / b_v \rceil\) (where \(b_v\) is the value-dimension tile size). On the Qwen3.5-2B shape with \(B_{\text{batch}}{=}1\), \(H{=}16\), \(b_v{=}32\) this yields \(64\) programs, well below the \(108\) SMs of an A100. Route II delegates the cross-chunk propagation to ChunkBwd, making each chunk’s fused kernel independent; the grid becomes \(N_C \cdot B_{\text{batch}} \cdot H \cdot \lceil d_v / b_v \rceil\) and at \(N_C{=}64\) reaches \(4096\) programs, a \(64\times\) increase that saturates SM occupancy. The shared cross-chunk ChunkBwd kernel itself runs the same \(64\)-program \((B_{\text{batch}}\!\cdot\!H\!\cdot\!\lceil d_v/b_v \rceil)\) grid in both routes and in the AR baseline.

7.2.7.3 HBM hand-off tensors.

Route I’s backward runs three kernels in strict sequence; each waits for the previous kernel’s output tensor to be fully written to HBM before launching. The two internal hand-off tensors have sizes \(L/B \cdot H \cdot d_k \cdot d_v\) (noisy-init gradients) and \(N_C \cdot H \cdot d_k \cdot d_v\) (chunk-boundary gradients); the former again carries the \(L/B\) factor. Route II collapses the first two stages and emits only the chunk-boundary tensor \(\mathbf{dh}^{\text{inject}}_{[c]}\) (\(N_C\)-sized) directly into ChunkBwd, halving the hand-off count and removing the \(L/B\)-sized intermediate.

7.2.7.4 Kernel latency and peak memory benchmark.

Tables 4 and 5 report an empirical wall-clock and peak-memory comparison of Route I versus Route II for the GDR and the ShortConv kernels, respectively, across block sizes \(B \in \{1, 2, 4, 8, 16, 32\}\) on the Qwen3.5-2B training shape. On the GDR, Route II cuts total wall-clock from \(135.1\) ms to \(37.7\) ms and peak memory from \(18.14\) GiB to \(456\) MiB at \(B{=}1\); Route I overtakes Route II at \(B{\ge}16\) once the chunk-level matmul saturates tensor cores. For the ShortConv, Route II is faster than Route I at every \(B\) in the sweep and holds peak memory at \({\sim}294\) MiB throughout. We therefore auto-dispatch Route II for the GDR at \(B{<}16\) and Route I at \(B{\ge}16\), and use Route II for the ShortConv at every block size.

4pt

7.3 1D Causal ShortConv: implementation routes↩︎

The 1D Causal ShortConv of width \(W\) (typically \(W{=}4\)) is the second sub-component of a GDN layer. It applies a depthwise causal 1D convolution to each channel of the input stream before the linear projections that feed the GDR of §7.2, conditioning each token’s \((\mathbf{q}_t, \mathbf{k}_t, \mathbf{v}_t)\) on its local temporal context. Under FLARE’s two-stream training, ShortConv is subject to the same class of constraints analyzed for the GDR in §7.1–§7.2.7: each noisy output position must read from the clean stream whenever its \(W\)-wide receptive field extends across the block boundary, while the clean stream is unaffected by the noisy stream. This subsection applies the GDR construction to ShortConv, identifies the two structural simplifications that reduce its cost, and states the two implementation routes in the same Route I / Route II terminology.

7.3.1 Two-stream contract for a width-\(W\) causal 1D convolution↩︎

Let \(\mathbf{w} \in \mathbb{R}^{D \times W}\) be the depthwise filter applied by a standard single-stream causal convolution to a length-\(L\) input of \(D\)-dimensional tokens; we write \(w_i\) for its \(i\)-th lag slice (acting channel-wise). Under the two-stream contract, the clean-side output \(y^{\mathrm{c}}_t\) at position \(t\) is computed from the clean stream in the standard form: \[\label{eq:short-conv-clean} y^{\mathrm{c}}_t \;=\; \sum_{i=0}^{W-1} w_i\,x^{\mathrm{c}}_{t-i},\tag{23}\] where \(x^{\mathrm{c}}_{t-i} \in \mathbb{R}^D\) is the clean-stream input at lag \(i\) (with out-of-range entries treated as zero) and \(w_i\,x^{\mathrm{c}}_{t-i}\) denotes the depthwise (channel-wise) product. The noisy-side output \(\tilde{y}_\ell\) at position \(\ell\) inside block \(b\) (positions \((b{-}1)B+1, \ldots, bB\)) reads from the noisy stream for lags that remain within block \(b\) and from the clean stream for lags that fall before block \(b\)’s start: \[\label{eq:short-conv-noisy} \tilde{y}_\ell \;=\;\sum_{i=0}^{W-1} w_i\,z_{\ell,i}, \qquad z_{\ell,i} \;=\; \begin{cases} \tilde{x}_{\ell-i} & \text{if } \ell-i \ge (b{-}1)B+1, \\ x^{\mathrm{c}}_{\ell-i} & \text{otherwise}, \end{cases}\tag{24}\] where \(z_{\ell,i} \in \mathbb{R}^D\) is the lag-\(i\) input read by the noisy output at position \(\ell\), selected from the noisy stream when the lag position falls inside block \(b\) and from the clean stream otherwise. The condition \(\ell - i \ge (b{-}1)B+1\) is the ShortConv analog of the clean-to-noisy visibility rule stated in Challenge 1 (§7.1): noisy tokens read from the clean stream only across the block boundary. Figure 24 illustrates the rule for \(W{=}4\).

7.3.2 Analysis of implementation difficulty compared to Gated Delta Rule↩︎

Two structural simplifications reduce the cost of the ShortConv case relative to the GDR case.

7.3.2.1 The noisy initial state is a raw tensor slice.

The GDR’s block-boundary clean state is a \(d_k \times d_v\) matrix obtained by running \(B\) per-step recurrence updates (Eq. 13 ) starting from \(\mathbf{S}_{[c]}\), incurring a serial replay or checkpoint lookup (§7.2.5, §7.2.6). The ShortConv seed required by Eq. 24 is the last \(W{-}1\) clean tokens preceding block \(b\)’s start, \(\bigl(x^{\mathrm{c}}_{(b{-}1)B-(W-1)+1}, \ldots, x^{\mathrm{c}}_{(b{-}1)B}\bigr)\), which is a slice of the clean-stream input tensor. No recurrence is replayed and no checkpoint tensor is constructed; the seed is read directly from the clean activations already consumed by the clean-side convolution.

7.3.2.2 The backward has no block-level state recursion.

The GDR’s backward (Fig. 14) composes four levels: a block-end readout (Level 1, Eq. 19 ), a per-step reverse sweep within a block (Level 2, Eq. 18 ), a block-level reverse recursion within a chunk through \(f^{\text{blk}}_b\) (Level 3, Eq. 17 ), and a cross-chunk matmul scan on chunk-boundary states (Level 4, Eq. 20 ). The ShortConv backward omits the block-level state recursion (Level 3): the gradient deposited by each noisy output into its clean-side lag positions is scattered directly into the corresponding positions of \(\mathbf{dx}\), without a block-level state-gradient recursion, because the noisy-side seed is not a recurrence terminus. The cross-chunk Level 4 reduces to the standard clean-side causal-conv backward and requires no additional hand-off.

7.3.3 Route I: batched initial-state↩︎

The ShortConv analog of Route I materializes the noisy initial states in HBM and then invokes the stock single-stream causal-conv kernel in batch. Forward:

Backward mirrors this structure. The batched noisy backward emits \(\mathbf{d}\tilde{\mathbf{x}}\), the filter/bias gradients, and a per-block gradient \(\mathbf{dH}^{\mathrm{init}}[b{-}1]\) with respect to each prepended context. A final scatter-add sends each \(\mathbf{dH}^{\mathrm{init}}[b{-}1]\) back to the last \(W{-}1\) positions of \(\mathbf{dx}\) that it was read from, and the clean-side causal-conv backward is run once on the full clean stream. The state tensor \(\mathbf{H}^{\mathrm{init}}\) lives in HBM throughout; its size is \(L/B \cdot D \cdot (W-1)\) bytes (times precision), the ShortConv analog of Route I’s \(L/B\)-scaling term for the GDR.

7.3.4 Route II: fused two-stream↩︎

Route II replaces both the init-state construction and the batched noisy conv with a single fused Triton kernel that implements Eq. 24 in place.

A companion fused kernel implements the backward: for each noisy output position \(\ell\) and lag \(i\), it computes the contribution \(w_i\,\partial\mathcal{L}/\partial \tilde{y}_\ell\) and atomically adds it to \(\mathbf{d}\tilde{\mathbf{x}}_{\ell-i}\) if \(\ell-i \ge \ell_{\mathrm{blk}}\), or to \(\mathbf{dx}_{\ell-i}\) otherwise; filter and bias gradients are accumulated into per-tile partial tensors and reduced. The clean-side output and its gradient are produced by the stock single-stream causal-conv forward and backward unchanged from the AR baseline. No \(L/B\)-scaling HBM tensor is materialized and no batched noisy-conv kernel is launched.

7.3.5 Implementation comparison↩︎

The three axes of §7.2.7 reapply in a reordered ranking, since the ShortConv state is small relative to the GDR state. In the deployed FLARE configuration (\(W{=}4\) with silu on CUDA, small \(B\)), Route II is preferable on all three axes; outside this configuration Route I remains a fallback.

7.3.5.1 Per-chunk parallelism.

Route I’s batched noisy conv has grid \((L/B) \cdot \lceil D/b_d \rceil\), where \(b_d\) is the channel-dimension tile size; Route II’s fused kernel instead tiles along the \(t\)-dimension, with grid \(N_t \cdot \lceil D/b_d \rceil\) for \(N_t\) the \(t\)-dimension chunk count, decoupling the launchable-program count from the block size \(B\).

7.3.5.2 Kernel-fusion overhead.

Route I issues three kernel calls per direction (forward: clean conv, init-state build, batched noisy conv; backward: batched noisy backward, init-state scatter, clean conv backward); Route II issues two (clean conv and fused two-stream conv, per direction). The launch-overhead savings are material at the short sequence lengths typical of ShortConv layers.

7.3.5.3 Peak memory.

Route I’s \(\mathbf{H}^{\mathrm{init}}\) has size \(L/B \cdot D \cdot (W-1)\), a few hundred MB per layer per sample at the smallest block sizes and shrinking with \(B\). This is much smaller than the analogous \(L/B\)-scaling tensor in the GDR case (gigabytes), so memory is not the dominant factor in selecting Route II for the ShortConv.

7.3.5.4 Kernel latency and peak memory benchmark.

Table 5 reports the empirical wall-clock and peak-memory comparison of ShortConv Route I versus Route II at \(B \in \{1, 2, 4, 8, 16, 32\}\) on the Qwen3.5-2B ShortConv shape. Route II is faster than Route I at every block size in the sweep and keeps peak memory within \({\sim}294\) MiB throughout, so FLARE uses Route II for the ShortConv at every block size.

4pt

7.3.5.5 Summary.

We auto-select Route II for every benchmarked ShortConv configuration supported by the fused kernel (\(W \le 4\) with silu/swish or no activation on CUDA), and retain Route I as a fallback for configurations outside it (e.g.alternative activations or larger filter widths).

7.4 Document-packed training↩︎

This subsection specifies the document-packing guarantees required by Challenge 2 of §7.1. The kernels in §7.2–§7.3 are described as operating on a single \(L\)-token document; for training throughput, multiple shorter documents are packed into each training sequence with per-document start offsets \(\mathbf{cu} = (c_0, c_1, \ldots, c_{N_{\text{doc}}})\), \(c_0 = 0\), \(c_{N_{\text{doc}}} = L\) (the standard cu_seqlens layout). The two-stream mask is defined per-document; under packing, the block-diffusion visibility rule is composed with a document-level mask that forbids cross-document attention, so each document is trained as if in isolation. This subsection states how the GDR and 1D Causal ShortConv kernels of §7.2–§7.3 preserve single-document semantics under packing.

7.4.0.1 Constraint: document starts are block-aligned.

Every entry of \(\mathbf{cu}[:-1]\) must be a multiple of the diffusion-block size \(B\), i.e.every document’s first token must also be a block start. This is a hard precondition: the two-stream block mask of §3.2 is defined on block boundaries, so a document starting in the interior of a block would cause its clean-to-noisy visibility to straddle a document boundary, violating the doc-level visibility contract. The final packed tail (\(c_{N_{\text{doc}}}\)) is permitted to be a partial block; the implementation zero-pads the trailing tokens up to a block boundary. A single helper enforces the check at every two-stream entry point; violations abort the training step with an explicit error.

7.4.0.2 Three kernel-level guards for cross-document correctness.

Under packing the three two-stream kernels of §7.2 and §7.3 each carry one document-boundary guard that the single-document versions do not need.

• Noisy recurrence’s initial state at document starts (Route II forward kernel of §7.2.6; same boundary also handled by Route I’s fill-in kernel). The fused two-stream kernel normally seeds the noisy recurrence of block \(b\) from the clean state \(\mathbf{S}_{(b-1)B}\) (Eq. 16 ). When block \(b\) is the first block of a document, \(\mathbf{S}_{(b-1)B}\) would be read from the previous document’s tail chunk and carry its clean-prefix information across the document boundary. The kernel detects this case (via the per-program is_doc_first_chunk flag built from \(\mathbf{cu}\)) and zeroes the noisy initial state instead. The forward’s noisy block \(b{=}1\) of every document therefore starts from \(\mathbf{0}\), matching the single-document semantics where each document is trained independently.

• Cross-chunk gradient shift at document boundaries (Route I fill-in backward of §7.2.5; also analogous code path inside the clean-side cross-chunk scan of §7.2.3.3). The normal cross-chunk shift turns \(\partial\mathcal{L}/\partial\mathbf{S}_{[c]}\) (a chunk-boundary gradient) into the previous chunk’s boundary gradient \(\mathbf{dh}_{[c-1]}\) for the scan. At a document boundary this shift must not cross: the boundary gradient of a new document’s first chunk has no corresponding end of a preceding chunk in the same document to flow into. The document-aware shift implements this by masking the cross-document assignment to zero; concretely, the contribution \(\mathbf{dh}_{[c-1]} \mathrel{+{=}} \mathbf{dh}_{[c]}\) is taken if chunks \(c-1\) and \(c\) are in the same document, and zeroed otherwise. Per-document chunk counts and offsets are computed once from \(\mathbf{cu}\) at the start of the backward step.

• ShortConv boundary read of the first noisy block of each document (Route II fused ShortConv kernel of §7.3.4). The lag-read rule of Eq. 24 lets a noisy output position read the preceding clean block when its receptive field straddles the block boundary. At a document boundary this read would cross into the previous document’s clean tokens. The kernel avoids this by bounds-checking every read with src >= 0 and src < doc_T after offsetting by the document’s bos; reads outside the current document contribute zero. For Route I the analogous behavior is enforced by zeroing the initial_states tensor at every document start, and by masking the corresponding \(\mathbf{dh}_0\) gradient in the backward so that no gradient flows back into the preceding document’s trailing clean tokens.

7.4.0.3 Variable-length kernel support.

In addition to the three guards above, every kernel in §7.2–§7.3 accepts a cu_seqlens argument with a compile-time IS_VARLEN switch. When cu_seqlens is provided, each Triton program loads its per-document bos/eos offsets from \(\mathbf{cu}\) and uses them to translate an in-program position into the correct global token index, replacing the default \(i_n \cdot T\) layout. This is standard variable-length kernel infrastructure inherited from the FLA reference implementation and is not FLARE-specific; packed and unpacked inputs share the same kernel code path.

7.4.0.4 Consequence for the cost model.

The quantitative claims of §7.2.7 and §7.3.5 (HBM scaling, backward-grid size, hand-off count) are stated in terms of a packed sequence length \(L = c_{N_{\text{doc}}}\); document count does not appear because the document-level guards above add \(\mathcal{O}(1)\) work per document and no additional per-layer tensors. In particular, Route II’s strided-checkpoint tensor has size \(N_C \cdot (M/S) \cdot H \cdot d_k \cdot d_v\) (bytes \(\times\) precision) regardless of how many documents are packed into the \(N_C\) chunks.

7.5 End-to-end MFU ablation↩︎

§5.2 reports the aggregate MFU lift produced by our kernel stack on FLARE-2B. This subsection breaks that number into its four incremental contributors and reports peak HBM alongside MFU at four block sizes, so that each MFU gain is attributable to a specific kernel change.

7.5.0.1 Setup.

All runs train FLARE-2B on a single node of \(8{\times}\)A100 80 GB in bf16, using fully-sharded data parallelism (FSDP) and no tensor-, pipeline-, context-, or expert-parallelism. Attention runs under FlexAttention; to fit A100 shared memory at head dimension \(256\), we instruct it to reduce the kernel’s pipeline depth to two stages. All runs use per-operator activation checkpointing (the PyTorch selective-AC policy with granularity “op”), a fused cross-entropy kernel, and the two training-side variance reduction options of §3.2: complementary-mask block pairs and antithetic mask sampling. The sequence length is \(L{=}4096\), the global batch is \(8{\times}\) the local batch (\(8\) for the local-batch-\(1\) rows, \(64\) for the local-batch-\(8\) rows), the training corpus is the Nemotron post-training v2 SFT mix (see Appendix 9), and the GDR chunk size (§7.2.1) is \(C{=}64\). Each row reports MFU and peak per-GPU HBM as the mean over measurement steps \(11\)–\(30\) of a \(30\)-step warm-up-then-measure run. Local batch is the per-GPU batch \(B_{\text{batch}}\); block size is the diffusion-block size \(B\) of §3.2. The AR reference is a plain Qwen3.5-2B run on the same tokens under a pure-AR single-stream model spec; it uses local batch \(1\) because that spec computes cross-entropy over the full logits tensor (the fused cross-entropy kernel above is only wired into the FLARE two-stream specs).

3pt

7.5.0.2 Per-row interpretation.

(1) Baseline. Both linear-attention sub-components run the Chunk-then-Refine path of §7.2.5. At \(B{=}1\) the row is out of memory; at \(B{=}16\) the baseline reaches \(22.06\%\) because the \(L/B\)-scaling state tensor is \(16{\times}\) smaller and the chunkwise matmul achieves high tensor-core utilization. (2) GDR Route II. Switching the GDR to the fused two-stream kernel of §7.2.6 unblocks \(B{=}1\) (\(12.56\%\)) and caps peak memory at every block size to \(12.7\) GiB, because the \(L/B\) state tensor is no longer materialized in HBM. At \(B{=}4\) the fused path is faster than Route I and lifts MFU by roughly \(+4\) points (and by \(\sim{+}1.8\) at \(B{=}8\)); at \(B{=}16\) it is slower, matching the crossover of §7.2.7 where the dense chunk-level matmul overtakes the fused kernel. (3) Raising the per-GPU batch from \(1\) to \(8\). The \(\sim 6\) GiB saved at \(B{=}4\) by setting (2) is spent on a larger per-GPU batch. No kernel code changes here; the memory headroom is converted into tensor-core utilization. The step is worth \(\sim{+}3.5\)–\(3.8\) MFU points at \(B{\in}\{4,8,16\}\) and pushes peak HBM to \(\sim 46\) GiB, still well within the \(80\) GiB budget. (4) ShortConv Route II (FLARE). The ShortConv fused kernel replaces \(L/B\) small per-block launches by a single fused kernel per sequence. The gain is largest at small \(B\) (\(+10\) points at \(B{=}1\)), since launch overhead scales as \(L/B\); it tapers off as \(B\) grows. The final row reaches \(24.81\%\) MFU at \(B{=}4\), matching and slightly exceeding the pure-AR reference of \(24.04\%\) on identical tokens. This is the ceiling one expects for a hybrid-backbone dLLM: FLARE processes both the clean and the noisy streams, and the complementary-mask training of §3.2 trains each block under both the masked and the unmasked partition in a single forward, so FLARE performs roughly \(4\times\) the per-input-token attention work of the AR baseline; matching AR MFU therefore means the kernel stack has absorbed that extra work without dropping hardware efficiency.

7.5.0.3 The \(B{=}16\) crossover is consistent with the microbenchmark.

Reading Tab. 6 at fixed \(B{=}16\), Route II for the GDR (row 2) underperforms Route I (row 1), for the same reason quantified in Tab. 4: at larger block sizes Route I’s dense chunk-level matmul overtakes Route II’s smaller launch footprint. Setting (4) still matches or beats setting (1) at \(B{=}16\) because the ShortConv Route II gain is additive and independent of the GDR choice.

7.5.0.4 Caveats.

Two properties of the measurement window should be kept in mind when comparing to single-stream MFU numbers reported elsewhere. (i) The complementary-mask training of §3.2 trains each block under both the masked and the unmasked partition in a single forward, doubling the effective batch seen by attention; turning this off would approximately double all FLARE MFU numbers in the table but would no longer reflect the production training regime of the released FLARE checkpoints. (ii) The FlexAttention pipeline-depth reduction above costs roughly \(1\)–\(2\%\) of attention throughput uniformly across all four settings and the AR reference, so it is already folded into every row. Absolute numbers on H100 or in FP8 will differ; the relative gains between settings should nonetheless be representative of the kernel stack’s contribution.

8 Efficient Inference and Implementation↩︎

8.1 Overview↩︎

Section 3.3 defines the two decoding interfaces used by FLARE, summarized in Figure 27. This appendix does not redefine those interfaces; instead, it records the serving-side details needed to implement them on a hybrid-attention backbone. We first show the path-specific masks and block layouts, then describe the proposal policies, fused verification kernels, denoising-loop controls, and recurrent-state machinery used by our SGLang-based serving stack [9].

8.1.0.1 Notation.

We write the prompt as \(x_{1:n}\), the generation buffer as \(y \in \mathcal{V}^{L}\) over vocabulary \(\mathcal{V}\) (initially \(y_i = \texttt{MASK}\) for \(i > n\)), the diffusion block size as \(B\), the per-block denoise step budget as \(S_{\mathrm{step}}\), and the top-\(1\) confidence \(c_i = \max_{v\in\mathcal{V}} p_\theta(y_i = v \mid \cdot)\) under the checkpoint’s output distribution \(p_\theta\). A confidence threshold \(\gamma \in (0,1)\), temperature \(T\), and truncation parameters \((k, p)\) parametrize sampling. On the AR-Trust side \(N\) is the speculation horizon, \(\mathbf{d} = (d_0, \ldots, d_{N-2})\) are the held draft tokens, and \(z\), \(\tilde{z}\) denote the clean-stream and noisy-stream logits emitted by one forward over the block (with \(z_i\), \(\tilde{z}_i\) their rows at position \(i\)); \(p_i\), \(q_i\) are the target (clean-stream) and draft (noisy-stream) distributions at the \(i\)-th verify position, formed from \(z_i\), \(\tilde{z}_i\) under the policy in use. We write \(\mathcal{B}_b = \{n + bB + 1, \ldots, n + (b{+}1)B\}\) for the position set of the \(b\)-th generation block. On the recurrent side, \(\mathbf{S}\) is the GDN state (the same \(\mathbf{S}\) as in Appendix 7), \(L_{\mathrm{gdn}}\) the number of GDN layers in the backbone, \(d_{\mathrm{state}}\) the per-layer state dimension, \(B_{\text{batch}}\) the scheduler batch size, and \(\mathbf{S}^{(r)}\) the state after \(r\) accepted updates inside a verify round. \(\mathcal{K}\) denotes a request’s KV cache slot set; writing a token \(y_i\) to \(\mathcal{K}\) means appending its key and value tensors as new entries. The set of GDN states indexed by \(\mathcal{K}\) is denoted \(\mathbf{S}[\mathcal{K}]\).

8.2 AR-Trust Sampling↩︎

Algorithm 5 gives the method-level interface. Here we specify the serving block layout and proposal policies used to realize that interface. Figure 28 (a) shows the attention pattern: the active block contains one anchor position (the last accepted token \(y_{\mathrm{anc}}\)), \(K \le N{-}1\) verify positions that hold the draft tokens to be accepted or rejected, and \(N{-}1\) MASK positions that produce the next round’s drafts. The anchor and verify positions are token-causal to the prefix and to each other; the draft positions attend bidirectionally over the whole \(2N{-}1\)-wide active block and causally to the prefix.

8.2.0.1 Block layout.

The active block has length \(2N{-}1\) (the maximum-width, steady-state shape); shorter rounds left-pad with MASK. In a cold-start round the block is \([t_0, \texttt{MASK}, \ldots, \texttt{MASK}]\) with \(1 + 2(N{-}1) = 2N{-}1\) positions: position \(0\) is a fresh clean token and the \(2(N{-}1)\) masked positions produce the next round’s \(N{-}1\) draft tokens (the first \(N{-}1\) MASK positions are unused padding on the cold start). In a verify round the block is \([y_{\mathrm{anc}}, d_0, d_1, \ldots, d_{K-1}, \texttt{MASK}, \ldots, \texttt{MASK}]\) of length \(1 + K + (N{-}1)\) where \(y_{\mathrm{anc}}\) is the last accepted token and \(d_0, \ldots, d_{K-1}\) are the \(K \le N{-}1\) draft tokens held from the previous round; at steady state \(K = N{-}1\) and the block reaches its full width \(2N{-}1\). The same forward returns clean logits at positions \(1, \ldots, K\), used to accept or reject each draft token, and fresh clean-plus-draft proposals from the \(N{-}1\) masked tail positions. Algorithm 29 expands the compact interface of Algorithm 5 with serving state, KV-cache bookkeeping, and recurrent-state updates.

8.2.0.2 Draft policies.

We consider three draft policies: Exact-Truncated, Softmax-Argmax, and Truncated-Argmax. They differ in how they form the proposal distribution \(q_i\) and, consequently, in the modified rejection rule and correction distribution they induce. The key subtlety is proposal consistency. Classical speculative decoding is distribution preserving only when \(q_i\) is the conditional law that actually produced \(d_i\) under the same proposal factorization assumed by verification. AR-Trust violates this condition if one simply recomputes \(q_i\) from current noisy-stream logits, because the noisy stream drafts multiple positions in parallel under a masked block context rather than an autoregressive proposal context. Exact-Truncated therefore stores the proposal probabilities used at draft time, while the two argmax policies intentionally approximate this condition to reduce serving overhead. Empirically, we do not observe significant quality degradation, but we separate these policies to make the exactness–efficiency trade-off explicit.

Exact-Truncated. Draft tokens are sampled from the same truncated distribution \(q_i\) (top-\(k\) and top-\(p\) applied to \(\tilde{z}_i / T\)) that verification will consult. The compact support \((\mathrm{top}\text{-}k\text{ ids}, \text{probs})\) of \(q_i\) is stored alongside each \(d_i\) until the next round. The corresponding modified rejection-sampling rule is given in Eq. 10 ; the stored \(q_i\) is the proposal distribution used when \(d_i\) was sampled, which is the closest match to the classical speculative-decoding condition.

Softmax-Argmax. Draft tokens are the argmax of \(\tilde{z}_i\), so \(q_i\) is a point mass at \(\arg\max \tilde{z}_i\). The ratio \(p_i(d_i) / q_i(d_i)\) degenerates and the accept rule reduces to a target-only check, with \(p_i^{\mathrm{full}}\) the full-softmax target over \(z_i / T\): \[\label{eq:accept-softargmax} \text{accept } d_i \iff u \leq \min\!\Big(1,\, p_i^{\mathrm{full}}(d_i)\Big), \quad y^\star \sim p_i^{\mathrm{full}} \setminus \{d_i\} \text{ (via Gumbel-max on } z_i/T \text{)}.\tag{25}\] No draft distribution is stored.

Truncated-Argmax. Draft tokens are the argmax of \(\tilde{z}_i\), but the accept rule uses the truncated \(p_i^{\mathrm{trunc}}\) and \(q_i^{\mathrm{trunc}}\) obtained by applying \((k, p)\) truncation to both \(z_i\) and \(\tilde{z}_i\): \[\label{eq:accept-truncargmax} \text{accept } d_i \iff u \leq \min\!\Big(1,\, p_i^{\mathrm{trunc}}(d_i) / q_i^{\mathrm{trunc}}(d_i)\Big), \quad y^\star \sim \mathrm{normalize}\!\big(\max(p_i^{\mathrm{trunc}} - q_i^{\mathrm{trunc}},\, 0)\big).\tag{26}\] Because \(d_i\) is obtained by argmax rather than by sampling from \(q_i^{\mathrm{trunc}}\), the induced distribution over accepted tokens does not match the target distribution exactly. This policy does not require storing \(q_i\) across rounds, and the verify kernel operates on a \(k\)-wide support.

8.2.0.3 Fused full-softmax verify kernel.

For the Softmax-Argmax policy, verification requires \(p_i^{\mathrm{full}}(d_i)\) on every verify row. A baseline implementation launches four or more kernels per row (softmax of \(z_i\), gather of \(p_i(d_i)\), the accept check, and, on rejection, construction and sampling of the correction distribution), and it materializes an \(\mathcal{O}(K V)\) probability tensor for the correction step at vocabularies of \(V \approx 1.5\!\times\!10^5\). This tensor dominates the AR-Trust decode step’s memory traffic: without fusion, verify becomes I/O-bound. We implement the verify step as a single Triton kernel that does not materialize this tensor. The kernel streams over \(z_i\) in tiles of \(\mathcal{O}(10^3)\) tokens, maintaining a numerically-stable running \((\max, \mathrm{sum\text{-}exp})\) pair in the style of FlashAttention [97], and computes \(p_i^{\mathrm{full}}(d_i)\) directly from the final pair and the logit at \(d_i\). Accepted rows exit the kernel after the accept check without performing correction-side work. Rejected rows take a second streaming pass and select the correction token by Gumbel-max over \(\log\max(p_i - q_i, 0) + \mathrm{Gumbel}(0, 1)\), tracking a running argmax across tiles. Gumbel-max samples the correction token without normalizing or CDF-inverting the correction distribution, removing two kernel launches that a separate CDF-inverse sampler would require.

8.2.0.4 Fused sparse verify kernel.

For the Exact-Truncated and Truncated-Argmax policies, \(p_i\) and \(q_i\) both have compact supports of size \(k\) (the top-\(k\) of each side’s truncation). The sparse verify kernel operates on \([k, k]\) tiles: loading the \(k\) token-id / probability pairs for target and draft takes \(4k\) floats plus \(2k\) integer indices per row. The correction \(\max(p_i - q_i, 0)\) is evaluated on the \(k\) token ids where \(p_i\) has mass via a \(k\!\times\!k\) lookup that reads \(q_i(v)\) for each \(v \in \mathrm{supp}(p_i)\), and inverse-CDF sampling over \(k\) entries selects the correction token. The per-row memory traffic is \(\mathcal{O}(k)\) instead of \(\mathcal{O}(V)\), which reduces the verify-kernel memory footprint by roughly three orders of magnitude at \(k \sim 50\) and \(V \sim 1.5 \times 10^5\); the kernel contains no runtime branches on the compile-time-known \(k\).

8.2.0.5 Tiled top-\(k\) logits over the vocabulary projection.

The draft policies require, at each masked position, the top-\(k\) ids and probabilities of the noisy-stream distribution. Let \(\mathbf{H} \in \mathbb{R}^{M \times d_{\mathrm{model}}}\) denote the final-layer hidden states at the \(M\) positions for which logits are needed, and \(W_{\mathrm{lm}} \in \mathbb{R}^{V \times d_{\mathrm{model}}}\) the LM-head projection. A baseline implementation materializes the full \(\mathbf{H} W_{\mathrm{lm}}^{\top}\) logits tensor of shape \([M, V]\) and then applies a top-\(k\). We tile the vocabulary-projection matrix multiplication: for vocabulary chunks of \(\mathcal{O}(3\!\times\!10^4)\) rows, we compute the partial \(\mathbf{H} W_{\mathrm{lm}}^{\top}\vert_{\text{chunk}}\) matrix, take its top-\(k\) per row, and merge the per-chunk top-\(k\) sets into a running top-\(k\) across chunks by concatenation and re-selection. Peak auxiliary memory is bounded by the size of the per-chunk partial matrix multiplication plus the running top-\(k\), and the full \([M, V]\) logits tensor is never instantiated. This is a chunked matrix multiplication with top-\(k\) merging rather than a fully fused Triton LM-head kernel; it avoids the dominant memory term in diffusion decoding and produces the compact logits that the sparse verify kernel consumes.

8.3 Diffusion-Trust Sampling↩︎

Algorithm 4 gives the method-level denoising interface. Here we describe the serving controls that make the same interface compatible with the hybrid state cache. Figure 28 (b) shows the attention pattern used during each denoising forward.

8.3.0.1 Per-block denoise loop.

A block begins with a clean seed at its first position followed by \(B{-}1\) masked positions. The seed is sampled from the final logit row of the previous block (or from the prompt prefill for block \(0\)). Each denoise iteration forms the unresolved set \(\mathcal{U}^{(s)}\), runs one bidirectional forward over the block, samples a candidate token at every unresolved position from the shifted noisy-stream logits under \((T, k, p)\), and commits the selected set \(\mathcal{C}^{(s)}\) according to the threshold schedule \(\gamma_s\). When all positions are set or the step budget is reached, a token-causal forward writes \(\mathbf{S}\) back to the state pool, appends the block to \(\mathcal{K}\), and produces the final logit row from which the next block’s seed is sampled. Algorithm 30 expands the compact interface of Algorithm 4 with seed sampling, state-pool updates, and per-forward control flags.

8.3.0.2 Differences between the denoise pass and the state-update pass.

Each denoise iteration operates on speculative (masked) rows whose committed values are not yet known, so writing the resulting recurrent state back to the live pool would contaminate \(\mathbf{S}\) with intermediate speculative updates. We therefore process every block in two forwards with different control flow. (i) The denoise forward uses the bidirectional block mask of Figure 28 (b); the state-update forward uses a token-causal mask over the block. (ii) The denoise forward reads \(\mathbf{S}\) from the recurrent pool but does not write it back; the state-update forward writes the final per-layer state back to the pool. (iii) The denoise forward returns the full vocabulary distribution at every block position because multiple positions are sampled per iteration, while the state-update forward returns only the final logit row for seed sampling. Our implementation exposes these three axes as flags on the forward-batch descriptor that the hybrid-attention kernels read to select mask type, state-update rule, and logits granularity.

8.4 Serving-system machinery for the hybrid backbone↩︎

8.4.0.1 Recovering the recurrent state at the accept boundary.

On a pure-softmax Transformer backbone, accepting \(r\) of \(K\) speculative tokens is a KV tail-trim: retain the first \(r\) new entries of \(\mathcal{K}\) and release the rest. The Gated DeltaNet recurrent state \(\mathbf{S}\) has no such per-position structure: the verify forward advances \(\mathbf{S}\) through \(K\) rank-one updates, and only the state after the \(K\)-th update is retained in the live pool. Two mechanisms can recover the required post-accept state \(\mathbf{S}^{(r)}\) (Figure 31). Path (a) – replay – re-executes the first \(r\) accepted updates through every GDN layer after the verify decision, at a cost of \(L_{\text{gdn}} \cdot r\) additional recurrent updates (matrix multiplications). Path (b) – cache-and-scatter – records \(\mathbf{S}^{(t)}\) for every verify position \(t \in \{0, \ldots, K{-}1\}\) during the verify forward itself, at a cost of \(L_{\text{gdn}} \cdot K\) additional stores (memory writes, no floating-point work). Draft-position intermediates are not stored, since the accept offset \(r \le K\) can never land on a draft position. We adopt path (b), because the store cost is dominated by the recurrent kernel’s existing matrix multiplication and is substantially cheaper than path (a)’s extra matrix multiplication. Concretely, the cache-and-scatter writes are produced inside the same per-step loop that already holds \(\mathbf{S}\) in registers: at the start of step \(t < K\), the kernel writes the current state \(\mathbf{S} = \mathbf{S}^{(t)}\) to a buffer of shape \([L_{\text{gdn}}, B_{\text{batch}}, K, d_{\mathrm{state}}]\) and then performs the step-\(t\) update, with no additional kernel launch. Once the verify decision is known, a single fused Triton kernel reads the accepted offset \(r\) per request, gathers the corresponding state slice from the buffer, and writes it to \(\mathbf{S}[\mathcal{K}]\) for every GDN layer. The kernel is launched with one program per \((\text{request},\;\text{layer},\;\text{state-tile})\) triple, so the \(B_{\text{batch}} \cdot L_{\text{gdn}}\) copies execute in one launch with per-tile parallelism over the state dimension, and requests with \(r = 0\) early-exit without memory traffic. A baseline PyTorch implementation of the same scatter step requires five operations (validity mask, \(\mathrm{nonzero}\) on valid request indices, two advanced-index selects, and a per-layer scatter), each a separate launch serialized over the layer dimension; our single-launch design removes this overhead.

8.4.0.2 Native dLLM mask modes and the prefix-tile fast path.

The mask required by both sampling paths is causal over the prefix, causal for the clean rows of the active block, and bidirectional for the masked rows. FlashInfer’s generic custom mask interface expresses this mask as a dense boolean tensor of size \(\mathcal{O}(B(B + n))\) per request. Our implementation replaces this encoding and the per-score predicate it triggers with three mechanisms in FlashInfer’s prefill kernel. First, two new mask modes accept compact metadata \((B, \text{row\_type}_{1..B})\) in place of the dense boolean tensor: one mode assumes a uniform block layout across the batch, the other admits per-request row-type variation. Second, the kernel exploits the following prefix-tile visibility invariant: for any KV tile whose indices lie entirely in the prefix, every score is unmasked regardless of query row type, because a causal clean row sees every prefix token and a bidirectional masked row also sees every prefix token. Third, when the kernel detects that a KV tile ends before the prefix boundary (a single integer comparison per tile), it skips the per-score mask predicate and takes the causal fast path used by non-dLLM decode. The compact encoding plus the prefix-tile fast path reduce the dLLM mask’s per-score overhead to zero on prefix-only tiles, which constitute the majority of attention compute at long context.

8.4.0.3 Graph-replay eligibility under diffusion-LLM invariants.

Fixed-shape decode blocks are captured into CUDA graphs to amortize launch overhead. A standard shape-based eligibility check is, however, either too strict for the two sampling paths (disabling graph capture whenever any dLLM flag differs from the captured state) or too lax (silently replaying the wrong graph under a mask-mode or logits-mode mismatch); the latter case corrupts AR-Trust verification. Our replay-eligibility predicate matches four invariants at replay time: (i) block size, (ii) mask-mode metadata, (iii) recurrent-state update mode, and (iv) logits-output mode. The draft policies of §8.2 differ on invariant (iv): Softmax-Argmax consumes dense logits whereas Exact-Truncated and Truncated-Argmax consume top-\(k\) logits. We therefore capture one graph per draft policy and dispatch through the four-invariant predicate. This preserves graph-capture speedups on the steady-state decode loop and eliminates the silent correctness failure mode.

9 Transfer-data mixes: sources and curation pipeline↩︎

This appendix documents every public dataset used to build the four transfer-data mixes of §4 (Long-CoT, Short-CoT+Math, Long-CoT+Math, Long-CoT+Math+IF), the common curation pipeline we apply to those sources, and the per-mix filtering and weighting decisions. §9.1 catalogs the source datasets with their size, sample format, prompt origin, and generator. §9.2 describes the shared preprocessing steps (message-schema unification, reasoning-trace handling, token estimation, filtering). §9.3 walks through each of the four paper-facing mixes and shows which sources it is built from and how its components are weighted at training time.

9.1 Source datasets↩︎

All sources are public SFT corpora with a shared [{role, content}] chat schema (some carry an auxiliary reasoning_content field that we merge into the assistant message, see §9.2). Each paragraph below gives a one-sentence summary of what the dataset contains, the sample count, the prompt source, and the generator that produced the assistant responses; the Hugging Face URL is in a footnote. Sample counts, storage sizes, and token statistics in this subsection are as reported on the dataset cards at the time of download.

9.1.0.1 Llama-Nemotron-Post-Training-Dataset (SFT splits)¹.

A post-training SFT mix of long-chain-of-thought traces across four splits, approximately \(4.2\)M samples total (math, \(2.22\)M; code, \(953\)K; science, \(708\)K; chat, \(349\)K). Each sample carries a reasoning field set to on or off; on samples contain an assistant response with an explicit <think>…</think> trace baked into the content. Prompts are drawn from NVIDIA’s curated post-training prompt pool spanning math competitions, coding benchmarks, science question-answering, and chat. Assistant responses are generated by the Llama-Nemotron synthesis pipeline (primarily DeepSeek-R1 and variants). Storage: \(114\) GB.

9.1.0.2 Nemotron-Post-Training-Dataset-v2².

A general-purpose post-training SFT corpus spanning four English categories (stem, chat, math, code) and five multilingual categories (Japanese, German, French, Spanish, Italian); only the English splits (\(1{,}397{,}187\) samples) are used in this paper. Average CoT length is shorter than Llama-Nemotron’s long-CoT traces (roughly \(800\)–\(1{,}350\) tokens per assistant message, depending on category). Assistant responses have <think> reasoning baked into content directly rather than carried in a separate field. Synthetic data is generated with Qwen3 and DeepSeek-R1-0528. Storage: \(92\) GB.

9.1.0.3 Nemotron-Instruction-Following-Chat-v1³.

Instruction-following chat conversations (\(320\)K samples) with an explicit capability_target label that distinguishes instruction-following from structured-output samples. Average assistant message is around \(3.9\)K tokens. Reasoning is stored in a separate reasoning_content field. Prompts come from NVIDIA’s instruction-following prompt pool; responses are generated with DeepSeek-R1-family models. Storage: \(6.4\) GB.

9.1.0.4 Nemotron-SFT-Instruction-Following-Chat-v2⁴.

The v2 refresh of the instruction-following chat set, with broader prompt coverage and \(1.99\)M samples. Average assistant message is around \(1.2\)K tokens (shorter than v1). Same reasoning_content schema as v1. Storage: \(14.4\) GB.

9.1.0.5 Nemotron-Science-v1⁵.

Science / STEM reasoning conversations (\(226\)K samples, average assistant message around \(2.5\)K tokens). Topics include physics, chemistry, biology, and STEM-style problem solving. Reasoning is stored in reasoning_content. Generators include DeepSeek-R1 and Qwen3-family models. Storage: \(2.3\) GB.

9.1.0.6 Nemotron-Math-Proofs-v1⁶.

Math-proof conversations (\(925\)K samples, average assistant message around \(4.3\)K tokens) emphasising symbolic derivations and proof-style chain-of-thought. Reasoning is stored in reasoning_content. Storage: \(27\) GB.

9.1.0.7 Nemotron-SFT-Competitive-Programming-v2⁷.

Competitive programming problem/solution pairs. The schema varies across its four splits: an exercism split (\(79\)K samples) with no reasoning traces; a text_to_sql split (\(97\)K); and two competitive-coding splits in C++ (\(333\)K) and Python (\(337\)K), both with reasoning_content traces averaging \(30\)K–\(56\)K tokens. Weighted total: \(845\)K samples. Prompt sources span Exercism exercises, text-to-SQL benchmarks, and CodeForces-style competitive problems. Responses are generated by DeepSeek-R1 variants. Storage: \(90\) GB.

9.1.0.8 Nemotron-Cascade-1-SFT-Data (Stage-2 instruction-following only)⁸.

NVIDIA’s first-generation “cascade” SFT dataset; we only use its instruction-following subset (\(146\)K samples). Prompts are taken from the Tulu-3 SFT mixture⁹; assistant responses are generated by DeepSeek-R1-0528.

9.1.0.9 Nemotron-Cascade-2-SFT-Data (instruction-following subset)¹⁰.

NVIDIA’s second-generation cascade SFT dataset. We use its instruction-following split after applying NVIDIA’s own constraint-satisfaction filter, which selects \(362\)K samples from a raw pool of about \(820\)K. Assistant responses are generated by the Cascade-2 pipeline (GPT-OSS-120B).

9.2 Curation pipeline↩︎

Every source goes through the same preprocessing pipeline before it is combined into any of the four transfer mixes. The pipeline has four stages, implemented in our data-preparation scripts and applied uniformly across sources.

9.2.0.1 (1) Message-schema unification.

All samples are rewritten into a single [{role, content}] chat schema. Three conversion variants handle the three schema families we encounter in the sources above:

• Sources with a separate reasoning_content field (Chat-v1, Chat-v2, Science-v1, Math-Proofs-v1, Competitive-Programming-v2) have it merged into the assistant content wrapped in a <think>{reasoning}</think> block.

• Llama-Nemotron’s input/output schema is flattened into a messages list; the <think> trace is already baked into the output string, so no reasoning merge is needed.

• Post-Training-v2 already ships with messages and <think> tags inline; we only filter to the English categories and pass the rest through unchanged.

Samples without any reasoning trace are padded with an empty <think></think> block so that downstream prompt-template handling is uniform across sources. The goal of this unification is that a plain string-concatenation of the resulting [{role, content}] messages under Qwen3’s ChatML template yields the same token sequence that a full Qwen3/Qwen3.5 apply_chat_template call would produce on the upstream raw sample. This lets our data pipeline avoid instantiating a tokenizer: every source is rewritten once into the canonical messages form, and the trainer handles tokenisation.

9.2.0.2 (2) Category and source tagging.

Each sample is assigned a category label (one of math, code, stem, chat) and a source tag identifying the upstream dataset. These labels are preserved through all subsequent stages and let us compute per-category statistics and apply category-specific filtering.

9.2.0.3 (3) Token-count estimation.

Because the combined source pool has tens of millions of samples, tokenizing the full corpus with the Qwen3-1.7B tokenizer is prohibitively slow. We instead estimate per-sample token counts from character counts plus a per-source char-to-token ratio: we tokenize \(3{,}000\) randomly-sampled messages per source to compute that source’s median character-to-token ratio, and divide the per-sample character count by this ratio (with a fixed \(15\)-token template overhead). Empirical spot checks against full tokenisation put the estimation error under \(5\%\). We record both a total-token and an assistant-token estimate for each sample.

9.2.0.4 (4) Filters applied uniformly.

Two filters run on every mix regardless of its target length or domain. First, samples whose assistant message contains an <think> tag without a matching </think> (indicating a truncated reasoning trace in the upstream source) are dropped; this removes approximately \(45\)K samples from the Long-CoT source pool and \(6\)K from the Math source pool. Second, when multiple sources are concatenated, we deduplicate across the combined pool using a hash of the serialized message list; this removes approximately \(11\)K duplicates when the IF sources are merged.

9.2.0.5 Output format.

After preprocessing each source yields a curated table where every row carries: the serialized messages list, the category and source tags, the estimated total and assistant token counts, and the sample’s origin metadata. The four transfer mixes of §9.3 are built by selecting, filtering, and concatenating rows from these curated source tables.

9.3 The four transfer-data mixes↩︎

The four mixes combine the curated sources of §9.1 under per-mix filtering and weighting rules. Each mix is trained at the same \(4{,}096\)-token packed context with global batch size \(256\) for \(9{,}000\) optimiser steps on the Qwen3-1.7B seed (§4).

9.3.1 Long-CoT↩︎

Long-CoT targets the high-reasoning-depth regime by keeping only samples whose estimated total token count is at least \(3{,}000\). It draws from seven sources of §9.1: Llama-Nemotron-Post-Training (all four SFT splits), Post-Training-v2 (English splits only), Instruction-Following-Chat-v1, Instruction-Following-Chat-v2, Science-v1, Math-Proofs-v1, and Competitive-Programming-v2. The construction proceeds as follows:

1. Apply the common preprocessing pipeline (§9.2) to each source.

2. For efficiency, apply a cheap character-count prefilter at twice the target token threshold (since the character/token ratio is bounded below by approximately \(2\)).

3. Apply the final filter using the calibrated total-token estimate, keeping samples with at least \(3{,}000\) tokens.

4. Concatenate the filtered per-source tables and shuffle with a fixed seed.

Final stats: \(4{,}588{,}156\) samples total, broken down as math \(2.55\)M, code \(1.23\)M, chat \(635\)K, stem \(179\)K. Storage: roughly \(186\) GB. At the \(4{,}096\)-token packing cap used during training, this corresponds to approximately \(18.3\)B total tokens and \(16.8\)B assistant (supervised) tokens.

9.3.2 Short-CoT+Math↩︎

Short-CoT+Math pairs a short-reasoning general mix with a dedicated Math source pool, giving a contrast point to Long-CoT. The short-reasoning source is Post-Training-v2 (English splits) used without any length filter; average assistant traces in Post-Training-v2 are roughly \(1\)K tokens, an order of magnitude shorter than Llama-Nemotron’s long-CoT output. The Math source pool is built from the math split of Llama-Nemotron (restricted to reasoning=on, roughly \(2.22\)M samples) concatenated with all of Math-Proofs-v1 (roughly \(920\)K samples after the unclosed-<think> filter drops \(6\)K). No length filter is applied to either component. The two components are passed to the trainer with equal weights (see Table 7). The Math source pool alone contains \(3{,}140{,}111\) samples; at the \(4{,}096\)-token packing cap it contributes approximately \(11.4\)B total tokens.

9.3.3 Long-CoT+Math↩︎

Long-CoT+Math is a direct intersection of the previous two mixes: the Long-CoT source pool from §9.3.1 (all \(4.6\)M length-filtered samples) and the Math source pool from §9.3.2 (Llama-Nemotron math plus Math-Proofs-v1, \(3.1\)M samples total). No length filter is applied to the Math source pool in this mix; the two pools are passed to the trainer with equal weights. This isolates the effect of adding a math-only anchor on top of the Long-CoT source pool without changing the short-CoT baseline.

9.3.4 Long-CoT+Math+IF↩︎

Long-CoT+Math+IF adds a dedicated instruction-following IF source pool to the previous mix. The IF source pool is built from four sources and deduplicated across the merged pool:

1. Nemotron-Cascade-1-SFT-Data Stage-2 IF (\(146\)K samples), which uses Tulu-3 SFT prompts and DeepSeek-R1-0528 responses.

2. Nemotron-Cascade-2-SFT-Data IF (\(362\)K samples, selected by NVIDIA’s constraint filter from a raw pool of roughly \(820\)K); responses are generated by GPT-OSS-120B.

3. Instruction-Following-Chat-v1 samples tagged with the instruction_following capability target (\(78\)K), giving a direct IF-labeled subset of Chat-v1.

4. Instruction-Following-Chat-v1 structured-outputs samples (\(5\)K), which ask the model to generate JSON or XML conforming to a given schema.

After preprocessing and deduplicating across the four parts on a hash of the serialized message list (which removes about \(11\)K duplicates), the IF source pool contains \(591{,}718\) samples, approximately \(0.66\)B total tokens and \(0.53\)B assistant tokens. It is trained alongside the Long-CoT and Math source pools with the weights of Table 7.

9.4 Mixing weights, sampling, and packing↩︎

Table 7 lists the trainer-level mixing weights used for the four mixes. The rest of this subsection describes how those weights are applied during data loading and how samples are packed into fixed-length training sequences.

9.4.0.1 Source-local packing.

Each source dataset in the mix is wrapped in its own SFT data-loading iterator that produces already-packed \(4{,}096\)-token training sequences. Raw samples from that source are concatenated into a running buffer via greedy bin-packing: for each raw sample we apply Qwen3’s ChatML template to produce a pair \((\text{tokens}, \text{labels})\) with \(-100\) on user and system positions (so only assistant tokens contribute to the loss); a sample whose tokenized length exceeds \(4{,}096\) is truncated, while a sample that would merely overflow the current buffer is deferred (the current buffer is flushed with padding up to \(4{,}096\), and the deferred sample starts the next packed sequence). Each packed sequence carries per-token document IDs; the trainer uses these with a document-causal attention mask to prevent cross-document attention within a packed sequence. For the AR training variant, the first label of each post-boundary document in a packed sequence is additionally set to \(-100\), so the shifted-label AR loss never asks the last token of one document to predict the first token of the next.

9.4.0.2 Per-sequence multinomial mixing.

Because packing is source-local, the unit of draw from the mixture is one fully-packed \(4{,}096\)-token sequence. On each next call, a mixer component is selected from the multinomial distribution \((w_1, \ldots, w_k)\) of the current mix, and that component’s next packed sequence is yielded. The multinomial draw is with replacement over components, while each component iterates its own packed sequences without replacement within its own epoch; the realized per-batch composition therefore matches \((w_1, \ldots, w_k)\) only in expectation. If a component is exhausted mid-epoch its weight is zeroed and the remaining weights are re-normalized. One consequence of this design matters for interpretation: raw samples from different sources never co-occur in the same packed sequence, so no cross-source attention can happen even under the document-causal mask. The sampling RNG is seeded deterministically and persisted in the training checkpoint, so a resumed run produces the same sample stream as an uninterrupted run.

4pt

9.5 Automatic SFT data selection for AR-to-dLLM transfer↩︎

The four transfer mixes of §9.3 select data primarily by length and per-source weighting. An orthogonal question is whether automatic, instance-level quality selection on the same raw SFT pool improves AR-to-dLLM transfer quality. This subsection documents such a pipeline: it scores individual examples with an instruction-following-difficulty signal, applies per-source quality gates, and rebalances the result by domain via cluster-aware sampling, all without manual curation. The resulting mixture (which we call the Auto-selected mix below) is not used by any of the paper-facing FLARE runs, but its controlled comparison against an unfiltered baseline (under the same Qwen3-1.7B seed and training budget) informed our final decision to keep the Long-CoT+Math+IF mix lightly filtered.

9.5.0.1 Source pools.

The automatic selection pipeline starts from two public SFT pools that together approximate the long-CoT plus instruction-following coverage of §9.3: Dolci-Think-SFT-32B (\(2.25\)M raw rows, \(23.58\)B raw tokens; \(1.17\)M rows and \(3.02\)B tokens at \(\leq 8\)K context) and Nemotron-Cascade-2-SFT-Data (\(24.49\)M rows, \(162.94\)B tokens; \(19.06\)M rows and \(54.04\)B tokens at \(\leq 8\)K context). Together the union has \(26.74\)M rows and \(186.52\)B tokens raw, dropping to \(20.23\)M rows and approximately \(57\)B tokens after the \(8{,}192\)-token Qwen3 ChatML length filter. Cascade-2 overlaps in part with the Nemotron-Cascade-2-SFT-Data instruction-following subset already described in §9.1; here we use the un-filtered \(\leq 8\)K pool rather than NVIDIA’s constraint-satisfaction selection.

9.5.0.2 Six-step curation pipeline.

The pipeline applies the following six steps in order. Each step only sees examples that survived the previous one.

1. Structural validity. Drop malformed conversations: empty messages, missing roles, zero-length content, unmatched <think> tags, assistant turns that repeat the prompt verbatim, and dangling final user turns. Known Cascade-2 tool-message leaks are normalized in place.

2. Length filter. Tokenize with the Qwen3 ChatML template and drop examples with more than \(8{,}192\) tokens. Long reasoning traces are dropped rather than truncated.

3. Deduplication and evaluation decontamination. Remove near-duplicates with MinHash LSH over \(5\)-gram prompt shingles at Jaccard \(\geq 0.8\) for intra-source dedup, and drop rows with Jaccard \(\geq 0.5\) against the prompts of GSM8K, MATH-500, IFEval, HumanEval, MBPP, ARC-C, and GPQA. This step follows standard large-corpus deduplication and decontamination practice.

4. Quality scoring. Score each instruction-response pair \((x, y)\) with the Instruction-Following Difficulty (IFD) score, defined as \(\mathrm{IFD}(x, y) = \mathrm{PPL}(y \mid x) / \mathrm{PPL}(y)\), where \(\mathrm{PPL}(y \mid x)\) is the response perplexity under the scorer when conditioned on the prompt and \(\mathrm{PPL}(y)\) is the unconditional response perplexity. Lower IFD indicates that conditioning on the instruction substantially reduces response perplexity, which we interpret as stronger instruction–response alignment. We use Qwen3-1.7B-Base as a cheap weak scorer. We additionally record rank-normalized IFD, response PPL, and response length, drop extreme response-PPL outliers, and do not score very short responses.

5. Composite ranking and source-specific gating. Rank rows within each source by a composite of normalized IFD, length, and a diversity weight: \(0.5 \cdot (1 - \mathrm{IFD}_{\text{norm}}) + 0.3 \cdot \mathrm{length}_z + 0.2 \cdot \mathrm{diversity}_{\text{wt}}\). Apply source-specific keep rates (top \(70\%\) for cleaner math/science/IF sources, top \(50\%\) for Dolci Python algorithms, and top \(40\)–\(60\%\) for noisier chat/tool sources).

6. Cluster-balanced sampling. Embed instruction text with BAAI/bge-small-en-v1.5, run \(k\)-means within each domain bucket, and round-robin sample from clusters by composite score until each bucket budget is met. Each curated row carries a cluster_id and an inverse cluster-frequency weight.

9.5.0.3 Target and realized domain mixture.

The bucket budgets were chosen to be evaluation-driven rather than proportional to raw source size, biased toward math, code, science, and instruction following. Table 8 lists the planned target shares and the realized bucket counts after sampling. The final mixture contains \(1{,}627{,}706\) rows and approximately \(5\)B tokens, against an original design target of roughly \(6\)B curated tokens for three epochs (\(\sim 18\)B effective training tokens).

4pt

The deviation between target and realized share, especially the overshoot on precise_if and science and the undershoot on math_direct and code, reflects how many high-IFD examples each upstream source actually contributed once the cluster-balanced sampler attempted to fill the bucket budgets without collapsing diversity.

9.5.0.4 Controlled comparison with an unfiltered baseline.

To assess whether the additional curation effort translated into downstream quality, we ran a head-to-head SFT comparison on Qwen3-1.7B at the same training budget against an unfiltered baseline mixture used elsewhere in the project (referred to here as the “baseline” mix). Both mixes were trained on the same seed checkpoint with the same optimizer settings and evaluated under the same \(16\)K-token evaluation harness. Table 9 reports the per-benchmark scores together with the parameter-matched Qwen3-1.7B reference evaluated at \(8\)K context.

4pt

Under this evaluation setting, the auto-selected mixture improves GSM8K, HumanEval, MBPP, and ARC-C relative to the unfiltered baseline mixture, while regressing on MATH-500 and IFEval. The pattern suggests that the IFD-driven composite ranking and cluster-balanced sampler effectively concentrate budget on cleaner code and science instruction-following data, but that the realized bucket undershoot on math_direct (Table 8) translates into a measurable MATH-500 regression, and the heavy reweighting away from generic chat hurts IFEval despite the high precise_if share.

We did not pursue this automatic selection pipeline further for the paper-facing FLARE runs. Both SFT runs in Table 9 fall well short of the reference Qwen3-1.7B at \(8\)K, indicating that under our fixed training budget the bottleneck is not a shortage of high-IFD examples but the same source-distribution shift discussed in §5: continuing SFT on external long-CoT data drawn from non-Qwen generators shifts the output distribution away from the post-trained seed. The FLARE Long-CoT+Math+IF mix in §9.3 therefore keeps filtering minimal and relies on length-based selection plus per-source mixture weighting rather than the heavier composite-ranking approach explored here.

10 Additional Experimental Details and Results↩︎

In this section, we present training and inference details for FLARE, as well as additional experimental results for the reader’s reference.

10.1 Inference and Training Hyperparameters↩︎

10.1.0.1 Training

For all experiments in this paper, we use an AdamW optimizer with a learning rate of \(10^{-5}\), a 500-step warm-up, and no learning rate decay. We always use a global batch size of 256 and a maximum sequence length of 4096 during training for 9000 optimizer steps. For Qwen3-related experiments, we use 32 NVIDIA RTX A100 80GB for experiments. For Qwen3.5-related experiments, we use 32 NVIDIA RTX H100 80GB for experiments. Each experiment requires a wall-clock time of 12 to 36 hours.

10.1.0.2 Inference

We use SGLang for all the benchmark inference. We use a max response length of 32768 tokens, sampling temperature 1, top-p 0.95, top-k 50 for all the model checkpoints.

10.2 Additional Experimental Results↩︎

10.2.1 Throughput of FLARE↩︎

14pt

Table 11 reports the full-family (2B/4B/9B) per-benchmark comparison of the released Qwen3.5 checkpoints, AR-SFT on the Long-CoT+Math+IF mix, and FLARE on the same mix, providing the source-model and AR-SFT references behind the main-text capability discussion in Section 5.

2.5pt

10.2.2 Study Summary Results↩︎

5pt

6pt

5pt

10.2.3 Per-benchmark Detailed Results of Main Tables↩︎

In this section, we include the per-benchmark score of each model we trained that is included in the main study and depicted in Tables 13–14. Every non-AR checkpoint is decoded under the AR-Trust sampling path through speculative decoding unless it’s non-applicable. mode (clean-stream trust path); AR baselines use native AR decoding. Benchmark columns group by capability (Knowledge & Instruction Following, Math, Code).

3.2pt

2.2pt

References↩︎

---

1. https://huggingface.co/datasets/nvidia/Llama-Nemotron-Post-Training-Dataset↩︎

2. https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2↩︎

3. https://huggingface.co/datasets/nvidia/Nemotron-Instruction-Following-Chat-v1↩︎

4. https://huggingface.co/datasets/nvidia/Nemotron-SFT-Instruction-Following-Chat-v2↩︎

5. https://huggingface.co/datasets/nvidia/Nemotron-Science-v1↩︎

6. https://huggingface.co/datasets/nvidia/Nemotron-Math-Proofs-v1↩︎

7. https://huggingface.co/datasets/nvidia/Nemotron-SFT-Competitive-Programming-v2↩︎

8. https://huggingface.co/datasets/nvidia/Nemotron-Cascade-1-SFT-Data↩︎

9. https://huggingface.co/datasets/allenai/tulu-3-sft-mixture↩︎

10. https://huggingface.co/datasets/nvidia/Nemotron-Cascade-2-SFT-Data↩︎