SGLang and Miles Add Day-0 Support for Inkling, a Frontier Multimodal Model

We're excited to partner with the Thinking Machines team to bring Day-0 support for Inkling to SGLang and Miles, with dedicated optimizations for its new architecture and broad feature coverage, and with Modal to train a DFlash draft model for Inkling.

Highlights

• Inkling is a 975B-parameter multimodal model with a context window of up to 1M tokens, featuring short convolution, attention with relative positional embedding, and a novel MoE design with a shared-expert sink.
• SGLang delivers up to 71.7k tok/s of input throughput and 171.0 tok/s of per-user decode speed on Nvidia Blackwell GPU, powered by dedicated optimizations for Inkling's architecture.
• SGLang provides broad feature support for Inkling, including speculative decoding, PD disaggregation, NVIDIA GPU support with bf16 + NVFP4 checkpoints, AMD GPU support with bf16 checkpoint, multi-LoRA serving, and HiCache.
• SGLang supports speculative decoding with DFlash, using a draft model trained by Modal specifically for Inkling.
• Miles implements Inkling in a customized Megatron backend with DP/PP/TP/SP/EP/CP support, enabling full-parameter and LoRA RL for text and vision-language tasks.
• Miles ensures train–inference consistency through customized kernels, routing replay, and cross-runtime parameter synchronization, while delivering steady RL improvements with both LoRA and full-parameter training across text-only and multimodal reasoning tasks.

SGLang Launch Commands: Cookbook

Miles Launch Commands: Documentation · Miles PR: miles#1683

Inkling Model Architecture

Inkling modifies the standard decoder-only Transformer architecture by integrating three components: short convolution, attention with relative positional embedding, and a shared-expert sink in MoE.

Overview

Figure 1: The model architecture of Inkling.

The figure above shows Inkling's designs for attention and MoE, each ending the same new way: a short convolution (ShortConv) before the residual connection. Attention also short-convolves K and V ahead of Q/K normalization, and adds a learned query-conditioned relative-position bias (RelLogitsProj) directly into the attention logits in place of positional embeddings. MoE runs top-k routing with a shared-expert sink (Sink). Unlike conventional shared-expert MoE — which adds an always-on shared path on top, independent of the router — Inkling's router scores the shared experts too and normalizes their weights together with the selected routed experts', so both draw from one shared weight budget.

Short Convolution (ShortConv)

Short Convolution is a short, per-channel causal convolution over the token dimension. For a hidden state at position $t$ and channel $c$, it reads the current token and the previous $W - 1$ positions of that same channel.

$$y_{t,c} = x_{t,c} + \sum_{d=0}^{W-1} w_{c,d}\, x_{t-d,\,c}$$

Current Inkling configurations use $W = 4$. ShortConv appears in four places per decoder layer:

• K stream ShortConv — right after the K projection, before Q/K normalization.
• V stream ShortConv — right after the V projection.
• Attention-output ShortConv — on the attention's output.
• MLP/MoE-output ShortConv — on the MLP/MoE's output.

Attention with Relative Positional Embedding

In Inkling, each attention layer adds a learned, per-head relative-position bias directly to the pre-softmax logits. Alongside $q_{t,h}$, $k_{t,h}$, and $v_{t,h}$, every query token carries a relative feature $r_{t,h}$. It is projected onto a vector of relative logits indexed by causal distance bucket $d$ (from $0$ up to $\text{rel\_extent}-1$), using a learned projection $P$:

$$\text{rel\_logits}(t,h,d) = r_{t,h}^{\top} P_{:,d}$$

The attention logit for query $i$, key $j$, and head $h$ adds a relative-position bias directly to the scaled dot product, selected from the query's relative logits by relative distance:

$$\ell(i,j,h) = \alpha \cdot q_{i,h}^{\top} k_{j,h} + \mathbb{1}[0 \le i-j < \text{rel\_extent}] \cdot \text{rel\_logits}(i,h,\,i-j)$$

where $\alpha$ is the attention scale factor. Causal masking handles future positions as usual; distances beyond $\text{rel\_extent}$ simply contribute zero bias in full attention.

Attention layout. Inkling mixes sliding-window and full-attention layers in the same stack — both variants use relative attention, just with different masking. The default layout repeats five sliding-window layers followed by one full-attention layer ($\text{layer\_id} \bmod 6 = 5$), giving the model cheap local layers with periodic full-context layers for long-range information. Full-attention layers additionally support a length-dependent log-scaling factor on $q_{t,h}$ and on the relative logits, letting the model adjust logit magnitude as context grows; local layers don't need this since their window size is fixed.

MoE With Shared Expert Sink

Inkling's feed-forward modules are standard sigmoid-gated top-k MoE — router scores, top-k selection, renormalized weights over the selected experts — with one exception: how its two shared experts are folded in. Most shared-expert MoE designs add an independent dense path on top, always-on experts that never compete with the router for probability mass. Inkling instead scores routed and shared experts together: top-k selection still picks only from the routed pool, but once that selection is made, the chosen routed scores and the shared-expert scores are concatenated and renormalized as a single group:

$$(w_{\text{routed}}, w_{\text{shared}}) = \operatorname{normalize}\big(\log \sigma(s_{\text{routed}} \,\Vert\, s_{\text{shared}})\big)$$

Both weight groups then scale their experts' outputs and sum into the same result.

SGLang Optimization

Inkling's three architectural changes each break an assumption a standard decoder-only server is tuned for. Making Inkling fast on SGLang meant new kernels and execution strategies for each.

ShortConv Optimization

ShortConv is small on its own — a $W=4$ per-channel causal filter, described above in Short Convolution — but it appears twice inside attention (K, V) and twice around the residual stream (attention-output, MLP-output), across every one of Inkling's decoder layers.

KV ShortConv + QK-norm + KV store. K and V stream ShortConv sit right before Q/K normalization and the KV-cache write. SGLang fuses all three steps — the K/V convolution, the Q/K RMSNorm, and the KV-cache write — into a single attention-prologue kernel.

All-reduce + ShortConv (+ RMSNorm and residual for decode). Attention-output and MLP-output ShortConv sit on the residual stream, right where the tensor-parallel all-reduce already runs — and all-reduce itself is expensive. SGLang built a family of custom all-reduce kernels to cut that cost, beating torch's multimem_all_reduce_ by up to 2.1×. Since ShortConv sits immediately after the all-reduce, SGLang fuses it directly into the all-reduce kernel instead of launching it separately. For decode, RMSNorm and the residual add are folded into the same kernel too, since decode's single-token shape makes the extra steps cheap. The fused kernel is 2.08–3.60× faster than the unfused chain, for +5–8% end-to-end throughput across the input-length sweep.

Prefill Full CUDA Graph. Breakable CUDA graph (BCG) and piecewise CUDA graph (PCG) are the standard way to cut CPU launch overhead during prefill: capture most of the forward pass once, then replay it without re-dispatching Python each step. Both still have to fall back to eager execution wherever an op needs live per-request metadata, and Inkling has four ShortConv sites per layer doing exactly that, on top of attention's own. Each site's eager fallback is a real, multi-level Python call stack, not a single kernel launch:

Figure 2: Under BCG, each of Inkling's four ShortConv sites (K stream, V stream, attention-output, MLP/MoE-output) falls back to a real, multi-level Python/Triton call stack instead of a single kernel launch — the eager CPU work that leaves a CUDA graph with bubbles.

With that much eager work per layer across roughly 66 layers, BCG and PCG still leave real CUDA bubbles — GPU idle time while the CPU works through the uncaptured segments — most visible at small-batch, moderate-context shapes where CPU dispatch competes with the GPU kernels for time.

Full CUDA graph capture removes those bubbles by capturing the entire forward, ShortConv included, into one graph sized to a capped number of request slots (full_prefill_max_req), so replay only refreshes buffer contents instead of re-dispatching any Python. Full-graph prefill throughput is roughly even with BCG at large shapes and +14–17% ahead of it at launch-bound ones.

Sharded ShortConv. Attention-output and MLP-output ShortConv run on the residual stream, so without sharding, every tensor-parallel rank redundantly computes and caches them at full width. SGLang instead offers a sharded strategy: reduce-scatter the partial sums so each rank holds only its hidden-dim shard, run ShortConv locally against a correspondingly sharded cache, then all-gather the result back to full width. This frees GPU memory, but the reduce-scatter/all-gather round trip is extra communication on top of the plain all-reduce, which costs more than it saves at decode's small per-step size — so it's an optional strategy rather than the default.

Relative Logits Optimization

Sheared-bias kernel. The relative bias term needs to be added inside the attention kernel itself. Inkling's FlashAttention-4 integration supports two equivalent approaches: a score-mod path, where each attention tile computes $i - j$ in its score callback and looks up $\text{rel\_logits}$ on the fly; and a sheared-bias path, where the relative logits are pre-sheared into a column-aligned bias tensor ahead of time,

$$\text{sheared\_bias}(i,h,j) = \text{rel\_logits}(i,h,\,i-j)$$

so the kernel can add bias via ordinary tile loads instead of per-score indexing. The sheared layout is faster but requires an Inkling-specialized FA4 fork with matching tile-alignment and padding conventions, and it's the path Inkling ships with. It's also the one place in the attention stack whose kernel schedule is cached by a buffer's identity rather than its contents, so anything that reuses that buffer's memory across calls — CUDA-graph replay included — has to make sure the cache gets refreshed rather than silently reused.

Figure 3: FA4's warp-specialized schedule in the MXFP8 + sheared-bias configuration. Amber marks the fork's additions: bias and scale-factor tiles staged alongside the existing Q/K/V loads, the bias add inside the online softmax, and the correction warps' V dequant.

Two-stream overlap. The rel_logits projection — turning $r$ into rel_logits via $P$ — depends only on the fused QKVR projection's output, not on the ShortConv/QK-norm work the attention-prologue kernel does to K and V. SGLang runs it on a separate CUDA stream from that prologue, so it overlaps with the conv/norm instead of serializing after it.

MXFP8 Support. Storing K/V in MXFP8 instead of bf16 roughly doubles KV-cache capacity, and Blackwell runs MXFP8 matmul natively at higher throughput than bf16. The FA4 fork applies this asymmetrically: Q@Kᵀ runs as a native MXFP8 MMA, while P@V stays in bf16 to protect accuracy, dequantizing V on the fly so the extra step overlaps the existing load/MMA schedule rather than costing a separate pass. Fusing the MXFP8 quantization directly into the same attention-prologue kernel, instead of a separate quantize-and-store step, keeps the cost to ~4.7–4.8 µs — within ~14% of the bf16 baseline — while carrying the KV-cache memory win.

Shared-Expert Sink Optimization

Top-k. The shared-sink MoE gate described above is sigmoid + biased top-k selection followed by a joint logsigmoid renormalization over the selected routed and shared scores — a sequence that, unfused, chains a sigmoid+bias op, a top-k op, and a renorm op. A single fused Triton kernel collapses that chain, 1.6–5.6× faster than the unfused chain across realistic token counts; a shape-specialized CUDA-JIT version goes further still, e.g. 7.72 µs vs the chain's 26.15 µs at T=4096 (3.4×) and 20.09 µs vs 52.88 µs at T=16384 (2.6×).

Shared-expert fusion. The shared-sink math (routed and shared scores sharing one normalization) means the shared experts run as one full-tensor-parallel dense block. Every token visits every shared expert, so preserving the expert axis as a batch dimension only creates extra work: the direct implementation replicates the input, runs batched GEMMs over the expert axis, materializes one output per expert, and sums those outputs.

We fuse that expert axis into the matrix dimensions instead. The gate/up weights are stacked along the output dimension, while the down weights are concatenated along the input (reduction) dimension. We call the resulting 2-D weight representation the linearized layout: “linearized” refers to folding the expert dimension into two dense GEMMs; the SwiGLU between them remains nonlinear. Its down GEMM performs the expert sum as part of its own reduction, eliminating replicated inputs, per-expert intermediates, and a separate sum.

Figure 4: Shared-expert fusion converts an expert-batched implementation into a linearized layout: gate/up weights are stacked into one dense GEMM, down weights are concatenated into a second dense GEMM, and the second GEMM's reduction performs the expert sum.

On B200 W4A16 serving, expert fusion lifts input throughput by 5.8–11.1% and cuts TTFT by 5.5–10.0% across BS1–32. H200 input throughput improves by 2.2–4.5%, while decode throughput holds steady.

Feature Support

Speculative Decoding with Multi-layer Eagle

Inkling ships with eight chained MTP (multi-token-prediction) layers for speculative decoding — not one draft head replayed eight times, but eight layers with their own weights, one per draft depth, each consuming the previous layer's output. Together they form an EAGLE-style recurrence: one target forward produces a hidden state, the chain extends it eight tokens deep, and the target verifies all nine positions in a single forward. The draft layers share the target's embedding and unembedding — SGLang aliases one copy across all of them.

Figure 5: One decode round of Inkling's multi-layer MTP. The target's hidden state and root token seed an eight-layer draft chain, captured in a single CUDA graph, and the nine-position block returns to the target for one verify forward.

One CUDA graph for the whole chain. The baseline design replays one graph per draft step, with Python between steps to sample a token, rotate it into the next step's input, and rebuild attention metadata — eight graph launches and seven host gaps per round. SGLang instead captures the entire chain in a single CUDA graph per batch size: all eight draft forwards, the token rotation, the per-step attention metadata, and the sampling itself, recorded back-to-back with no Python in between. This works because the chain is shape-invariant across steps and sampling is graph-safe, drawing fresh randomness on every replay; the draft window is fixed-width, so a request that accepted 2 tokens and one that accepted 7 replay the same graph. A single fused kernel feeds the graph from the verify output, replacing the ~20 small copies the host used to issue between verify and draft.

Distribution-exact rejection sampling. At temperature 0 the verify step accepts by matching the target's argmax. At temperature > 0, Inkling serving supports true speculative rejection sampling (--speculative-use-rejection-sampling): accept draft token $X_k$ with probability $\min(1, p_k(X_k)/q_k(X_k))$ and resample from the residual distribution on rejection, which preserves the target distribution exactly rather than approximating it with acceptance thresholds. Verify then needs each step's full draft distribution $q_k$, not just the sampled token — so the graphed chain stashes each step's probabilities into a persistent [bs, 8, vocab] buffer as it goes, and the verify kernel walks the tokens in chain order against the stash.

A sync-free decode round. All of the above is stitched together so the decode loop runs fully overlapped: across the entire round — draft chain, verify, and the glue between them — SGLang issues no device synchronization at all. The host stays ahead of the GPU, and kernels from consecutive phases stream back-to-back without the pipeline ever draining.

Speculative Decoding with DFlash

SGLang also supports speculative decoding with DFlash, a standalone draft architecture rather than layers chained onto the target. A DFlash draft runs as its own model with its own KV cache, filling in a full block of masked future positions in one forward instead of stepping through them autoregressively; it borrows the target's embedding and LM head rather than training its own, and the target verifies the whole block in a single linear (non-tree) forward. We collaborated with the Modal team to train a DFlash draft model for Inkling, giving Inkling serving a second speculative-decoding path (--speculative-algorithm DFLASH) alongside the native multi-layer MTP chain above.

Prefill-Decode Disaggregation

Inkling has three heterogeneous state types: full-attention KV, sliding-window KV, and ShortConv convolution state. All three need to be transferred from the prefill node to the decode node. The three-component pool was already built for hybrid SSM-attention models, and Inkling's convolution state fits the same infrastructure, so no new transfer logic was needed on the disaggregation path. SGLang's UnifiedRadixCache is the common abstraction over all state types; the disaggregation layer, HiCache, and speculative decoding all compose over three components through that single interface.

AMD GPU Support

SGLang supports Inkling on AMD MI35X through a generic score_mod interface added to the Triton attention backend. This interface inlines a caller-supplied score modification function into the extend and decode kernels at compile time, so Inkling's relative position embedding works on ROCm without the NVIDIA-only FlashAttention-4 path. Combined with additional operator adaptations and an aiter MoE runner integration, Inkling runs end-to-end on MI35X with --attention-backend triton --moe-runner-backend aiter.

LoRA Serving

For a projection $y = Wx$, LoRA leaves the base weight $W$ unchanged and adds a low-rank update:

$$y = Wx + \alpha B(Ax), \qquad r \ll d.$$

$A$ shrinks the token from the model dimension to rank $r$, and $B$ expands it back. Inkling applies these updates across attention, dense MLP, expert, and output projections, so executing every update after its base GEMM would turn many small low-rank operations into exposed latency.

SGLang instead uses a two-stream schedule. The main CUDA stream runs the base model, while a LoRA side stream computes the rank-$r$ work for the adapter selected by each token. The streams synchronize only where the delta is needed; fused join kernels expand and add the delta while applying the surrounding activation. In an MLP, for example, the gate/up shrink overlaps its base GEMM before the streams join at

$$\operatorname{SiLU}(g + \Delta g) \odot (u + \Delta u).$$

The down-projection base and low-rank paths then run toward a second join. Expert layers reuse the base model's routing metadata rather than routing the LoRA update again. Requests selecting different LoRAs stay in the same launches: a per-token adapter mapping selects the appropriate $A$ and $B$ factors without splitting the batch by adapter.

Figure 5: Inkling's two-stream LoRA schedule keeps base-model GEMMs on the main CUDA stream and overlaps the selected adapter's low-rank work on a side stream. Fused join kernels add the deltas at the activation and output dependencies.

The result stays close to the no-LoRA baseline even when a batch contains several distinct LoRAs. The table reports TP8 output throughput (input 8192 / output 1024, symmetric memory on, no speculative decoding); single-LoRA rows use --max-loras-per-batch 1, and 4-LoRA rows use --max-loras-per-batch 4 with 4 distinct adapters in the batch:

config	single-LoRA BS1	single-LoRA BS4	4-LoRA BS4
B200 W4A16 TP8	123.20	383.57	380.17

Moving from one distinct LoRA to four at BS4 costs 0.9% on B200.

Radix Cache and HiCache

Reusing KV cache across requests that share a prefix is one of the most effective optimizations in LLM serving, especially for multi-turn and long-context workloads. SGLang's radix cache captures it by organizing shared prefixes into a radix tree, and HiCache tiers that tree across memory (L1 GPU HBM, L2 host DRAM, L3 disk or a remote store) so a prefix evicted from the GPU reloads from a cheaper tier instead of being recomputed once the working set outgrows HBM. Inkling complicates both, because it does not present the single homogeneous attention KV pool that prefix caching assumes. Most of its layers use sliding-window attention (SWA), while its ShortConv branch keeps a small per-request convolution state that SGLang stores in the pool built for recurrent Mamba/SSM layers, even though Inkling has no SSM. SGLang has a separate radix cache for each axis: a full-attention RadixCache, a SWARadixCache for sliding-window eviction, and a MambaRadixCache for recurrent state. A hybrid model that exercises several axes at once needs them coordinated on one tree, which SGLang's UnifiedRadixCache (#20415) provides as a single tree over typed components (FULL, SWA, MAMBA) with HiCache supported natively on top rather than reimplemented per variant.

Inkling plugs straight into this path. It registers a three-component pool, and the unified cache selects a combined SWA + Mamba stack strategy that builds one host stack over all three components (full KV, SWA KV, and the convolution-state pool) behind a single cache controller. Because the integration is native to the unified cache rather than a bespoke side path, Inkling inherits HiCache's tiered reuse and composes with the rest of the serving stack, including speculative decoding.

Multimodal Optimization

SGLang offloads the compute-intensive portions of image preprocessing (patchification, normalization, and content hashing) to a native Rust extension that runs GIL-free within the serving process. This enables multi-image parallelism and moves hash computation off the scheduler hot path, reducing TTFT by 14-44% and improving serving throughput by ~20% for image-heavy workloads. The extension is built around an extensible processor interface and currently supports Inkling.

Performance Results

We benchmark Inkling end-to-end on a B200 node, sweeping batch size at a fixed sequence shape (input length 8192, output length 1024) under both TP=4 and TP=8, with symmetric memory and CUDA graphs enabled throughout.

Figure 7: Inkling serving performance on B200 at ISL=8192/OSL=1024. Left: throughput per GPU vs. interactivity, swept over batch size. Right: inter-token latency at matched batch sizes.

The left panel traces the throughput–interactivity Pareto frontier at both tensor-parallel sizes: moving along either curve trades per-user decode speed (interactivity) for aggregate throughput per GPU (total input and output tokens served, per GPU, per second), by admitting more concurrent requests. At TP=8, batch size 1 reaches 171.0 tok/s/user, and batch size 32 reaches 71.7k tok/s of aggregate input throughput. The right panel shows inter-token latency staying in the single-digit milliseconds through batch size 8 at both tensor-parallel sizes, only climbing into the tens of milliseconds once batch size reaches into the 30s and beyond.

Reinforcement Learning with Miles

Miles delivers Day-0 RL support for Inkling through a Megatron backend spanning full-parameter and LoRA optimization across DP/PP/TP/SP/EP/CP. To maintain train–inference consistency, the backend combines customized relative-attention, ShortConv, and FP32 MoE kernels with rollout routing replay, which replays shared-sink MoE routing over both text and media-expanded sequences. Building on this foundation, native LoRA covers Inkling's attention, dense MLP/MoE, and LM-head projections with adapter-only synchronization, while the multimodal pipeline supports image and audio inputs. We validate the complete RL pipeline with 975B full-parameter GRPO, text LoRA, and vision-language LoRA.

The complete Inkling RL implementation is available in Miles pull request #1683. A ready-to-use Inkling image is available with docker pull radixark/miles:inkling.

RL Training Backend with DP/PP/TP/SP/EP/CP Support

Miles implements Inkling as a native Megatron model. The backend reconstructs Inkling's local and global relative attention, four residual short-convolution paths, dense-to-MoE layer schedule, shared-sink router and experts, and image and audio encoders as differentiable Megatron modules. It supports all six parallel dimensions used by the Inkling training recipes:

• DP. Distributes microbatches and synchronizes gradients across replicas.
• PP. Partitions transformer layers across pipeline stages.
• TP. Shards the fused $q$/$k$/$v$/$r$ and output projections, with $k$/$v$ ShortConv aligned to local heads.
• SP. Shards residual ShortConv paths while preserving causal context.
• EP. Partitions routed experts and their trainable state.
• CP. Uses contiguous all-gather: $q$ and $r$ stay local, while $k$ and $v$ are gathered with global offsets for relative attention. ShortConv gathers the full sequence and returns each rank's local slice.

An Inkling model bridge handles checkpoint conversion in both directions: it shards the released Hugging Face weights for Megatron training, then reconstructs the layouts required by checkpoint export and rollout weight updates. The same backend drives both full-parameter and LoRA RL.

Full-Parameter RL

Built on the native Inkling backend and parallel stack above, Miles runs full-parameter GRPO by updating every model tensor. In each iteration, SGLang rollout workers generate trajectories and record the routed-expert IDs required by R3. Miles training ranks consume those trajectories, replay the routing decisions, and apply the full-model update. The model bridge then converts the updated distributed shards into SGLang's tensor layout and streams the next policy back in bounded buckets for the following rollout.

Within the bounded GPU memory capacity of a single GB300 rack, the full policy, gradients, and FP32 optimizer state create additional memory pressure. Miles therefore streams Megatron DistributedOptimizer state between a bounded GPU working set and node-local NVMe. This GPU–disk offload changes storage placement, not the optimizer update (miles#1575, torch_memory_saver#80, Megatron-LM#63).

We run Inkling 975B full-parameter GRPO on 12 nodes with 4 GB300 GPUs per node. Training uses DP2/PP3/TP4/EP8, while rollout uses TP8/EP16. The run uses a global batch size of 32, a GRPO group size of 8, and a maximum response length of 4K with truncation. It enables routing replay and uses Adam with a learning rate of $10^{-6}$. We train on DAPO-Math-17K and evaluate on AIME25. Train–rollout KL remains around $10^{-3}$, while both raw reward and AIME25 evaluation improve steadily.

Figure 8: Inkling 975B full-parameter GRPO: train–rollout KL remains around 10⁻³ while raw reward and AIME25 evaluation improve steadily.

Customized Efficient Kernels for Train–Inference Consistency

Train–inference consistency is a central correctness requirement in RL: the trainer must evaluate the same policy that generated the rollout, even though training and serving run through different distributed execution stacks. Maintaining that consistency is difficult because the same checkpoint can produce different token probabilities when kernels, accumulation precision, packed-sequence boundaries, or reduction orders differ. Miles addresses this mismatch at the operators that define Inkling's forward pass, while supplying the backward implementations needed for distributed training.

Relative attention. The fused projection produces the per-token, per-head features $q_{i,h}$, $k_{i,h}$, $v_{i,h}$, and $r_{i,h}$. Miles applies the same fixed relative projection $P$ and attention-logit definition used by SGLang:

$$\text{rel\_logits}(i,h,d)=r_{i,h}^{\top}P_{:,d}, \qquad \ell(i,j,h)=\alpha\,q_{i,h}^{\top}k_{j,h}+b(i,j,h).$$

Here $b(i,j,h)$ selects $\text{rel\_logits}(i,h,i-j)$ under the same causal-distance rule defined above.

Miles provides three training-side attention backends for Inkling: FlexAttention, FA4, and Transformer Engine, with FlexAttention used by default. Each forward first computes the compact relative logits $rP$. Miles implements the relative-score modification with a customized CUTE kernel, which indexes $rP$ by relative distance inside each attention tile and preserves SGLang's causal or sliding-window rule without materializing a dense token-by-token bias matrix. Miles prepares and caches the score modifier and block mask for each attention geometry, amortizing their construction across repeated packed-sequence shapes.

FlexAttention provides a differentiable forward and backward through both the attention scores and $r_{i,h}$, while $P$ remains frozen. At an 8K packed sequence and relative extent 1024 on GB300, it is roughly $5\times$ faster and uses roughly $5\times$ less peak memory than the Transformer Engine reference, while matching it within BF16 numerical noise. FA4 remains available as an optimized alternative and Transformer Engine as a reference backend.

Short convolution. Inkling applies residual causal ShortConv to the $k$ and $v$ streams, the attention output, and the MLP output. Miles implements this operator with customized Triton forward and backward kernels that fuse the depthwise causal convolution with the residual path. The kernels accumulate in FP32 and add the residual last, reproducing SGLang's arithmetic order before casting back to the model dtype. For packed sequences, Miles precomputes the start and end of each token's segment: the forward kernel masks left-context reads before the segment start, while the backward kernel masks gradient reads beyond the segment end, preventing state or gradients from crossing sample boundaries. Under sequence parallelism, Miles gathers the full sequence context before ShortConv and then returns the corresponding output and gradient shard to each rank.

FP32 MoE activation and combination. Inkling's shared-sink MoE is sensitive to rounding in both the gated activation and the weighted expert reduction. Miles executes these two stages in FP32. A customized differentiable Triton SwiGLU kernel performs the activation and per-token shared-expert scaling before a single cast back to the model dtype; expert outputs are then accumulated in FP32, following SGLang's summation order, and cast once after combination. This aligns the continuous expert computation across training and rollout, complementing the discrete routing alignment provided by R3.

Routing Replay for the Shared-Sink MoE

Beyond customized precision-aligned kernels, Rollout Routing Replay (R3) is equally important for train–inference consistency in MoE RL. A small numerical perturbation near a top-$k$ boundary can change the selected experts and therefore the computation graph itself. SGLang records the routed top-$k$ expert IDs during rollout, and Miles reuses those IDs in the training forward.

For token $t$, let $e_1,\ldots,e_k$ be the replayed routed experts. Miles recomputes their current router scores and places them before the scores of all always-active shared experts:

$$s_t=\left[ \sigma(\ell^r_{t,e_1}),\ldots,\sigma(\ell^r_{t,e_k}), \sigma(\ell^s_{t,1}),\ldots,\sigma(\ell^s_{t,N_s}) \right].$$

All entries then use one common normalization:

$$w_{t,m}=c\frac{s_{t,m}}{\sum_n s_{t,n}}.$$

The first $k$ weights belong to the replayed routed experts, and the remaining $N_s$ weights belong to the shared experts. Here $c$ combines Inkling's configured route scale and learned global scale. R3 replays only the expert IDs; Miles recomputes the continuous weights from the current router, so gradients still flow through both routed and shared weights. The implementation uses a softmax over log-sigmoid scores for numerical stability. Truncated importance sampling handles the remaining probability drift after the expert subgraph has been aligned.

The second Inkling-specific adaptation is multimodal indexing. SGLang records routing over the sequence after image patches and audio features have expanded the prompt, so Miles uses the engine-reported expanded length rather than the original text-token count. The trace is then padded and sliced with the same packed-sequence and sequence-parallel layout as the training batch before it is registered to each MoE layer.

LoRA Implementation in Inkling

Inkling's released LoRA schema covers attention, dense MLP, MoE, and LM-head projections. Miles implements the same structure natively in Megatron, including TP/EP-aware execution and direct export to SGLang, so RL updates only the adapter while the base model remains frozen.

For each adapted linear layer, the base weight $W$ stays frozen and Miles trains only the low-rank factors $A$ and $B$:

$$y=Wx+\frac{\alpha}{r}B(Ax).$$

The Inkling-specific challenge is applying this standard update efficiently across its heterogeneous attention, dense, and expert projections.

Inkling's routed experts use a shared-outer factorization instead of an independent pair of factors for every expert. For $w_1$ and $w_3$, $A$ is shared across experts while the expert-specific $B$ tensors follow EP sharding. For $w_2$, the expert-specific $A$ tensors are sharded and $B$ is shared.

Let $\mathcal C$ index all adapted projections. Their low-rank factors form the trainable adapter state

$$\Phi_t = (A_{\ell,t}, B_{\ell,t})_{\ell\in\mathcal C}.$$

With $\Theta_0$ denoting the frozen base model, the policy at iteration $t$ is

$$\pi_t(\cdot\mid x)=\pi(\cdot\mid x;\Theta_0,\Phi_t).$$

GRPO updates $\Phi_t$ to $\Phi_{t+1}$ while $\Theta_0$ remains fixed. Before the next rollout, Miles synchronizes only $\Phi_{t+1}$; the base weights remain resident in both runtimes.

End-to-End LoRA RL with Adapter-Only Synchronization

Figure 9 follows one complete LoRA RL iteration. Along the upper path, SGLang returns sampled tokens, rollout log-probabilities, rewards and masks, and the routed expert IDs required by R3. Miles evaluates the same token sequence, injects the recorded expert IDs at each router, and differentiates only the adapter. In MoE layers, the low-rank branches consume the already-dispatched expert-token buffers, reusing the base model's expert grouping instead of routing the LoRA update separately. Expert-specific factors follow EP sharding, while shared outer factors are replicated and their gradients are reduced across the EP group.

Figure 9: Colocated LoRA RL: SGLang sends rollout trajectories to Miles, while Miles returns only the updated adapter and keeps the frozen base resident.

After the optimizer step, Miles materializes a serving-ready adapter directly from the distributed training state. Instead of issuing one collective per tensor, the exporter packs the requested TP and EP fragments and performs one flat all-gather per parallel group. It then reconstructs Inkling's SGLang layout: fused projection shards are split and concatenated in serving order, expert tensors return to expert-major order, and padded LM-head rows are removed. The result is a set of contiguous BF16 tensors with the exact names and shapes consumed by SGLang.

The lower path performs the version switch without staging adapter payloads through host memory. Miles pauses generation, flushes the engine cache, serializes the named GPU tensors as CUDA IPC handles, and sends each payload to the colocated SGLang worker sharing that GPU. SGLang unloads the previous adapter, then loads the complete next version in a single call and shards it into the serving layout before generation resumes. The frozen base either remains resident on the rollout side or is synchronized once before iterative adapter updates; subsequent steps transfer only the adapter. Released Inkling adapters in safetensors format support warm starts, while native per-rank adapter checkpoints retain optimizer and scheduler state for exact training resume.

Using the same experimental setting as the full-parameter run, LoRA GRPO keeps train–rollout KL on the order of $10^{-3}$ while raw reward improves steadily over roughly 450 steps. Adapter-only synchronization reduces weight-update latency from 49.4 seconds to 2.5 seconds, a $20\times$ speedup, while restricting backward and optimizer work to the adapter reduces training-step time to 85% of full-parameter training.

Figure 10: Text LoRA GRPO: raw reward improves over roughly 450 steps.

Multimodal LoRA RL with Expansion-Aligned Routing Replay

Inkling is a powerful multimodal model released by Thinking Machines Lab, with native support for text, image, and audio inputs. Miles extends the RL backend across these modalities: both full-parameter and LoRA recipes accept structured multimodal rollouts, execute Inkling's vision and audio towers in Megatron, and preserve routing replay across the media-expanded sequence.

Figure 11 shows the central data transformation. Miles keeps the original structured message list until the Inkling-specific renderer emits the model's role and content markers together with exactly one sentinel for each image or audio item. The processor checks that the sentinel count matches the supplied media, then records $p$ image patches or $f$ audio d-mel frames for each item. Before training, Miles replaces one image sentinel with $p$ in-vocabulary placeholder positions and one audio sentinel with $f$ positions, while recording the sample-local positions at which the corresponding media embeddings must be inserted.

Figure 11: Each media sentinel expands into the model positions occupied by image patches or audio frames, and the R3 routing trace is indexed over the resulting expanded sequence.

SGLang makes its MoE routing decisions after this expansion, so the R3 trace contains expert IDs for every expanded position rather than only for the rendered text sequence. Miles validates the trace length against the text length plus the recorded media expansion, expands the token sequence before packing, and pads and slices R3 with the same layout used by the training tokens. Because media sentinels occur in the prompt, response log-probabilities and loss masks remain unchanged. During batching, the recorded sample-local media positions are shifted into the packed-global sequence; the Megatron vision and audio towers encode the patch and d-mel tensors and scatter their embeddings into exactly those positions. Under sequence parallelism, each rank selects the positions that fall in its local sequence shard.

In multimodal LoRA RL, the current production recipes keep the vision and audio towers frozen and train the language-model adapter or base model around their embeddings. Training the media towers is available as an experimental option.

To validate the multimodal RL recipe, we choose a vision-and-text setting on Geo3K, split the dataset into training and evaluation subsets, and keep the parallel configuration and optimization hyperparameters identical to the preceding LoRA experiment. The run exercises the same LoRA execution, routing replay, and adapter synchronization on visual math prompts. Geo3K evaluation rises from roughly 0.54 to 0.58 while train–rollout KL remains on the order of $10^{-3}$ throughout training. The same multimodal path also supports audio inputs, full-parameter training, and distributed execution.

Figure 12: Vision LoRA GRPO: Geo3K evaluation improves with training steps.

Acknowledgment

This work was a collaboration between the SGLang & Miles team and Thinking Machines Lab.

SGLang & Miles team: Ke Bao, Cheng Wan, Chunan Zeng, Zhichen Zeng, Yanbin Jiang, Yuhao Yang, Qiaolin Yu, Mao Cheng, Yi Sun, Mingyi Lu, Haoguang Cai, Banghua Zhu, Ying Sheng

Thinking Machines Lab: Aurick Qiao, Paul Zhang, Shenxiu Liu

Thanks to the Modal team for collaborating with us to train the DFlash draft model for Inkling.