VibeSys on Trainium: Agents Build/Tune Qwen3.8 Serving Engines on AWS Trainium2

Figure 1. VibeSys run for Qwen3.8-27B on one trn2.3xlarge.

In the VibeServe paper, we laid out a vision for agent-driven bespoke systems. Generic serving engines do well on common cases but miss optimization opportunities on the long tail. For that tail, coding agents make it affordable to generate a bespoke engine per target. Our AMD MI300A post tested this thesis on AMD GPUs.

We now extend the scope to AWS Trainium2, targeting Qwen3.8 models (Qwen3.8-27B and Qwen3.8-Flash-Next). Trainium2 is an even harder target because of its different programming model and thinner open-source support. On Trainium, the recommended serving path is vLLM Neuron, with its model recipes covering only Llama 3, gpt-oss, dense Qwen3-32B, and a few other Qwen3 variants, not supporting Qwen3.8 models.

Porting inference engines written for other hardware also doesn’t work. Trainium’s compiler needs static shapes and decides memory layouts itself, and its NeuronCores are programmed through NKI, not CUDA.

Qwen3.8’s architecture makes the serving system even harder to implement. Both models interleave Gated DeltaNet linear attention, which keeps a recurrent state per sequence, with gated full attention. Flash-Next adds QSA sparse attention with a learned indexer, a 512-expert MoE, hyper-connections, an n-gram embedding table, and an MTP draft layer. Each mechanism needs its own state, routing, or kernels.

Together, these models and this hardware form a combination that generic systems have not reached. This is exactly the kind of target where a bespoke system has the most to gain.

We use TorchNeuron, the native PyTorch backend that AWS released in October. It exposes Trainium2 as a neuron device. Code runs either in eager mode, where each operation is compiled the first time it sees a new shape, or through torch.compile(backend="neuron"), which compiles a whole captured graph into one executable with neuronx-cc. NKI kernels are called from PyTorch as custom operators.

To support Trainium in VibeSys, we modified the framework to provide three things:

Setup. Qwen3.8-27B has 64 layers (48 GDN, 16 attention) and 54 GB of BF16 weights. Serving it requires sharding the model across all four NeuronCores of a trn2.3xlarge, which has one Trainium2 chip and 96 GiB of HBM. The workload is 64 concurrent requests of 1k prompt tokens and 1k generated tokens, and the metric is output throughput.

The starting point is the Hugging Face Transformers implementation of the model, which is generic PyTorch code that also serves as the correctness reference but has few serving optimizations and nothing written for Trainium. From it, agents wrote the whole engine in a fresh repository, including the model layers in PyTorch and NKI, the scheduler, the caches, and the HTTP server. Existing serving engines and model libraries such as vLLM Neuron and NxD Inference were off limits. Every round had to pass an accuracy gate that compares greedy outputs and teacher-forced logprobs against Transformers on 12 prompts, including 4k-token and cache-resume cases. A candidate was accepted only if it beat the current best by more than the 2% noise bar.

All agents ran on Opus 5.5. The 37 rounds took about 53 hours and $230 in API usage.

Results. The first working engine served 195 tok/s, 50 minutes in. By round 3, it served 425 tok/s, and the final engine (round 37) serves 860.7 tok/s (Figure 1). That cut the benchmark’s wall time from 336 s to 76 s. A decode step for 64 sequences now takes 56 ms, down from 111 ms after round 3. The final engine is about 2,400 lines of Python, including three NKI kernels.

Fix the compiler’s state updates (224 → 425 tok/s). In round 3, the agents found that an indexed copy of one slot’s GDN state was slower than a dense masked where over all 64 slots (57 vs. 19 ms per prefill layer group). A one-pass decode recurrence also cut a decode layer group from 10.4 to 6.7 ms.

Stop moving the GDN state (425 → 584 tok/s). Decode moved about 1.1 GB of FP32 state per layer group, almost four times the minimum. A dense prefill inverse, double-buffered state, and a single-pass NKI decode kernel fixed this. The agents wrote the kernel by following the compiler’s instruction stream, so its output is bit-identical. VibeSys’s accuracy gate caught an earlier, faster but buggy kernel that flipped a near-tie token.

Stop re-laying out memory (584 → 778 tok/s). Compiled decode attention spent about 1.0 ms transposing the KV cache before 0.25 ms of real work, so a hand-written NKI attention kernel gave +27%. Pipelining decode with a non-blocking copy and an event wait, instead of .cpu(), which waits for the whole device queue, saved another 4 ms per step.

Trim prefill (778 → 861 tok/s). Smaller changes added 0.3–5% each.

Setup. We serve Qwen3.8-Flash-Next (48 layers, BF16) on a trn2.48xlarge with 16 Trainium2 chips. We optimize throughput for a burst of 64 requests of 8k input and 1k output tokens.

Running and debugging many experiments at this scale, with a 48-layer MoE model across 16 chips, is extremely time-consuming, and a single new kernel’s warmup can take close to an hour. So we started by supporting the model in ServingStudio, our in-house simulator for LLM serving systems (PR). That work had four parts.

With the simulator, agents could rank changes before spending hours on hardware.