KDFlow: Decoupled LLM Distillation Framework

• KDFlow is a decoupled training framework that separates teacher inference on SGLang from student training on PyTorch FSDP2 for efficient LLM distillation.
• It reduces communication overhead by transferring only hidden states and locally recomputing teacher logits, achieving a 37× reduction in data transfer.
• Experiments demonstrate that KDFlow matches state-of-the-art distillation quality while providing 1.44× to 6.36× speedups, especially with Mixture-of-Experts teachers.

Searching arXiv for papers on KDFlow and related usages. KDFlow is a distillation-focused training framework for LLMs that features a decoupled architecture and employs SGLang for teacher inference, thereby bridging the training efficiency of FSDP2 and the inference efficiency of SGLang in a unified system (Zhang et al., 2 Mar 2026). In adjacent literature, the label also appears in two other senses: as a description of how different knowledge components propagate through student gradients in DeepKD, and as a kernel-density-based flow viewpoint for Wasserstein gradient flows of KDE-approximated divergences in generative modeling (Huang et al., 21 May 2025, Cao et al., 11 Mar 2026). The term is therefore context-dependent: in LLM distillation it denotes a concrete systems framework, while in nearby work it names or describes structured information flows.

1. Terminological scope

Context	Meaning of “KDFlow”	Representative paper
LLM distillation systems	A decoupled teacher–student framework	(Zhang et al., 2 Mar 2026)
DeepKD context	Gradient-mediated knowledge flows in distillation	(Huang et al., 21 May 2025)
Generative modeling context	A kernel-density-based Wasserstein gradient flow viewpoint	(Cao et al., 11 Mar 2026)

In the LLM setting, KDFlow addresses a systems mismatch in standard knowledge distillation pipelines. Existing frameworks such as TRL, MS-SWIFT, EasyDistill, and OpenRLHF typically place both teacher and student inside the same training backend, even though the teacher only runs forward passes while the student runs forward, backward, and optimizer updates. KDFlow reassigns these roles to specialized backends instead of using a homogeneous training stack (Zhang et al., 2 Mar 2026).

A separate usage appears in DeepKD, where “KDFlow” is best understood as how different kinds of knowledge are injected into, and propagated through, the student via the gradients. DeepKD formalizes three such flows—Task-Oriented Gradient, Target-Class Gradient, and Non-Target-Class Gradient—and reshapes them through decoupled momentum and denoising (Huang et al., 21 May 2025). A third usage appears in Gradient Flow Drifting, where the terminology refers to a kernel-density-based transport flow in Wasserstein space, implemented through KDE score differences (Cao et al., 11 Mar 2026).

2. Decoupled architecture for LLM distillation

KDFlow’s central architectural move is to decouple the teacher and student into specialized execution environments. The teacher side runs on SGLang, an inference-optimized backend; the student side runs on PyTorch FSDP2, a training-oriented distributed primitive; and Ray serves as the process and distributed execution manager. The system is organized around a single Trainer process and three actor groups: TeacherActorGroup, StudentActorGroup, and RolloutActorGroup (Zhang et al., 2 Mar 2026).

The motivation is explicitly systems-oriented. Training engines are optimized for gradient computation, optimizer state and gradient sharding, checkpointing, and fault tolerance, but the teacher only requires high-throughput, low-latency inference. The mismatch becomes particularly acute for Mixture-of-Experts teachers, because training engines such as FSDP and ZeRO-3 do not handle sparse expert routing and efficient expert parallelism as effectively as specialized inference runtimes such as vLLM or SGLang. In the reported analysis, teacher forward time can dominate per-step training time when MoE teachers are run under FSDP-type backends (Zhang et al., 2 Mar 2026).

TeacherActorGroup embeds an SGLang engine instance directly rather than the HTTP server, which avoids HTTP overhead and improves ndarray compatibility. StudentActorGroup performs full student training under FSDP2, including forward, backward, and optimizer steps. RolloutActorGroup supports on-policy KD by running the student in inference mode using SGLang Router and SGLang servers, with CUDA Interprocess Communication used to update SGLang weights directly from the training student (Zhang et al., 2 Mar 2026).

This design is not a change to the distillation objective itself. Its significance is that it separates inference and optimization concerns while preserving white-box KD semantics.

3. Hidden-state transfer and exact distillation semantics

A core systems problem in decoupled KD is communication. If teacher logits are transmitted directly, the cost becomes prohibitive. For a typical LLM, teacher logits have shape $B \times L \times V$; for 128 sequences of length 4096 with the Qwen3 vocabulary $|V| = 151{,}936$, BF16 logits require

$$128 \times 4096 \times 151{,}936 \times 2\text{ bytes} \approx 160\text{ GB}.$$

KDFlow avoids this by transferring only the teacher’s final hidden states $h^T \in \mathbb{R}^{B \times L \times d}$ and recomputing teacher logits locally on the student side through the teacher LM head: $z^T = W_{\text{LM}^T} h^T + b_{\text{LM}^T}.$ Because typically $d \ll V$, the communication volume ratio becomes

$$\frac{B L d}{B L V} = \frac{d}{V} \approx \frac{4096}{151{,}936} \approx 0.027,$$

which corresponds to roughly a 37× reduction in transferred elements (Zhang et al., 2 Mar 2026).

The hidden states are produced as Numpy ndarrays inside the SGLang engine and passed through Ray’s Shared Object mechanism in shared memory, enabling zero-copy cross-process access. Student processes therefore receive a pointer into shared memory rather than a serialized payload. The student then combines its own forward pass with locally reconstructed teacher logits and applies the KD loss (Zhang et al., 2 Mar 2026).

KDFlow frames this mechanism as preserving exact KD losses rather than approximating them. In standard white-box KD, the student minimizes divergences between full teacher and student output distributions, for example forward KL: $$\mathcal{L}_{\text{FKL}} = \sum_t \mathrm{KL}\big(p_T(\cdot \mid x_{1:t}) \,\|\, p_S(\cdot \mid x_{1:t})\big),$$ with $p_T = \mathrm{softmax}(z^T/\tau)$ and $p_S = \mathrm{softmax}(z^S/\tau)$. Hidden-state transfer followed by local recomputation yields the same $|V| = 151{,}936$0 as direct teacher-side logit computation, assuming the same LM head weights and numerical precision. This is distinct from top-$|V| = 151{,}936$1 or partial-logit schemes that alter the KD objective (Zhang et al., 2 Mar 2026).

The main trade-off is compute for communication. KDFlow adds a student-side matrix multiplication for $|V| = 151{,}936$2, but the framework treats that cost as minor relative to the communication and memory burden of transmitting full logits.

4. Training regimes, KD objectives, and extensibility

KDFlow supports both off-policy and on-policy distillation. In off-policy KD, the Trainer samples prompt–response pairs $|V| = 151{,}936$3 from a fixed dataset, sends them to TeacherActorGroup for teacher hidden-state extraction, then forwards the batch and $|V| = 151{,}936$4 to StudentActorGroup for student training. In on-policy KD, RolloutActorGroup first generates responses $|V| = 151{,}936$5 from the current student policy through SGLang-hosted rollout servers; those prompt–response pairs are then scored by the teacher, and StudentActorGroup trains on the resulting data before weights are synchronized back to rollout actors through CUDA IPC (Zhang et al., 2 Mar 2026).

The framework exposes these regimes through dedicated trainer abstractions. Off-policy training uses OffPolicyKDTrainer, while on-policy training uses OnPolicyKDTrainer. The operational difference is confined to orchestration: static dataset iteration in the former, rollout–teacher–student cycling in the latter (Zhang et al., 2 Mar 2026).

KD objectives are also modular. Built-in losses include Forward KL, Reverse KL, Jensen–Shannon Divergence, and Total Variation Distance. Cross-tokenizer KD is supported natively through DSKDv2, which supplies a mapping mechanism between teacher token-level distributions in the teacher vocabulary and student token-level distributions in the student vocabulary, enabling a proper KD objective when tokenizers differ (Zhang et al., 2 Mar 2026).

The API is deliberately structured so that the systems pipeline is separated from the algorithmic layer. A user specifies the teacher model, student model, KD objective, distillation mode, and data pipeline, while Ray-based actor orchestration, hidden-state transfer, zero-copy communication, and rollout synchronization remain internal. A plausible implication is that experimentation on KD objectives becomes less coupled to distributed-systems implementation.

5. Experimental profile and comparative results

The reported experiments use teacher models Qwen3-14B, Qwen3-32B, and Qwen3-30B-A3B; student models Qwen3-4B and Qwen3-1.7B; 100k prompts sampled from LMSys-Chat-1M; responses generated with Qwen3-14B for off-policy experiments; and a single server with 8× NVIDIA H20 GPUs. Baselines include TRL with DeepSpeed ZeRO-3, MS-SWIFT with ZeRO-3, ROLL with FSDP2, and a pure FSDP implementation in which teacher and student run in the same FSDP process. The main metrics are training time per iteration, KL loss curves, and downstream performance on AlpacaEval 2.0 via LC-Win Rate and Win Rate (Zhang et al., 2 Mar 2026).

On loss correctness, KDFlow’s KL loss curve for distilling Qwen3-30B-A3B into Qwen3-4B essentially overlaps the pure FSDP KD baseline. The paper also reports that FP8 teacher inference yields almost identical curves to BF16, and interprets this as evidence that the decoupled design and SGLang-based teacher inference do not introduce significant numerical issues for KD (Zhang et al., 2 Mar 2026).

On downstream performance, the off-policy FKL experiment distilling Qwen3-30B-A3B into Qwen3-1.7B yields the following AlpacaEval 2.0 results: Qwen3-1.7B without KD reaches LC-Win Rate 26.09% and Win Rate 21.99%; MS-SWIFT reaches 28.40% and 27.86%; the pure FSDP baseline reaches 28.18% and 28.20%; and KDFlow reaches 28.23% and 28.32%. These numbers place KDFlow at essentially the same quality level as the strongest baselines while maintaining its systems advantage (Zhang et al., 2 Mar 2026).

On speed, the framework reports 1.44× to 6.36× speedup compared to current KD frameworks. The largest gains occur with MoE teachers. For the Qwen3-1.7B student and Qwen3-30B-A3B teacher, MS-SWIFT requires 36.9 seconds per iteration, ROLL requires 53.8, KDFlow BF16 requires 5.9, and KDFlow FP8 requires 5.8. For the Qwen3-4B student and the same teacher, MS-SWIFT requires 43.2 seconds per iteration, ROLL 67.9, KDFlow BF16 11.3, and KDFlow FP8 11.1. This suggests that the decoupled backend choice is especially consequential when sparse expert routing makes teacher inference the dominant bottleneck (Zhang et al., 2 Mar 2026).

6. Related uses of “KDFlow” in adjacent research

In DeepKD, “KDFlow” denotes the way knowledge enters the student through gradients, rather than a distributed framework. DeepKD decomposes the logit-level training signal into Task-Oriented Gradient, Target-Class Gradient, and Non-Target-Class Gradient, corresponding respectively to hard-label task supervision, teacher confidence on the true class, and teacher dark knowledge among non-target classes. It then allocates separate momentum buffers to these components and introduces a Dynamic Top-k Mask that gradually expands the set of non-target classes participating in distillation. The paper reports that KD+DeepKD improves ResNet32×4→ResNet8×4 on CIFAR-100 from 73.33% to 76.69%, and to 77.03% with top-$|V| = 151{,}936$6; on ImageNet, KD+DeepKD with top-$|V| = 151{,}936$7 improves ResNet50→MobileNet-V1 to 74.65% versus 70.50% for vanilla KD; and on MS-COCO, KD+DeepKD with DTM improves AP by approximately +1.93 over vanilla KD (Huang et al., 21 May 2025).

The significance of that usage is conceptual rather than infrastructural. DeepKD treats knowledge transfer as a multi-flow optimization problem, with GSNR used to motivate different momentum coefficients for TOG, TCG, and NCG. In that setting, “KDFlow” refers to gradient decomposition, temporal smoothing, and denoising rather than teacher–student systems design.

A different usage appears in Gradient Flow Drifting, where the term identifies a kernel-density-based flow on probability measures. The paper proves that the drifting field of the Drifting Model equals, up to a bandwidth-squared factor, the difference of KDE log-density gradients: $|V| = 151{,}936$8 This is presented as the particle velocity field of the Wasserstein-2 gradient flow of $|V| = 151{,}936$9 under KDE approximation. The same framework also includes MMD-based generators as special cases and proposes a mixed reverse KL plus $$128 \times 4096 \times 151{,}936 \times 2\text{ bytes} \approx 160\text{ GB}.$$0 strategy intended to avoid both mode collapse and mode blurring. The experiments are preliminary and limited to synthetic 2D benchmarks, where forward KL and $$128 \times 4096 \times 151{,}936 \times 2\text{ bytes} \approx 160\text{ GB}.$$1-type drifting show good mode coverage but can blur between modes, while the reverse KL plus $$128 \times 4096 \times 151{,}936 \times 2\text{ bytes} \approx 160\text{ GB}.$$2 mixture is described as “almost only precise” with little blur and still explores all modes quickly (Cao et al., 11 Mar 2026).

Taken together, these usages show that “KDFlow” is not a single canonical method. It can denote a concrete LLM distillation system, a gradient-structured view of classification distillation, or a KDE-based Wasserstein flow in generative modeling.

7. Limitations, assumptions, and open directions

KDFlow for LLM distillation is explicitly research-oriented rather than a fully industrial training stack. The student backend is currently limited to PyTorch FSDP2, and the paper notes that FSDP2 does not yet match the training efficiency of highly optimized industrial 3D-parallel systems such as Megatron-LM at very large scale. The framework also lacks some industrial optimizations, including asynchronous training, sophisticated scheduling across thousands of GPUs, and advanced failure recovery mechanisms (Zhang et al., 2 Mar 2026).

The framework is also dependent on SGLang as the inference engine. The reported experiments treat that choice as robust, including for FP8 teacher inference, but future work is identified around integration with other inference engines and further handling of numerical quirks, particularly for MoE or low-precision inference (Zhang et al., 2 Mar 2026).

Several misunderstandings recur around the term. One is that decoupling teacher and student necessarily implies approximating teacher supervision; KDFlow is explicitly designed to preserve exact KD losses by transferring hidden states and recomputing logits locally. Another is that the label names a single theory or software object across the literature; the available record instead shows a framework name in LLM distillation and two context-specific descriptive usages in distillation dynamics and generative modeling. This suggests that precise disambiguation is necessary whenever the term appears.

The forward-looking directions stated for the LLM framework are concrete: integrating Megatron-LM or other 3D-parallel training systems as alternative student backends, implementing asynchronous training and more sophisticated pipeline parallelism, extending cross-tokenizer KD support with more algorithms such as ULD and multi-level OT, and exploring deeper on-policy KD variants and hybrid KD+RL setups (Zhang et al., 2 Mar 2026).