This document is relevant for: Trn2, Trn3

High-performance GPT-OSS 120B implementation for 8K/1K serving on Trn3#

Starting with Neuron SDK 2.33, vLLM Neuron has a high-performance implementation of GPT-OSS 120B for 8K-input, 1K-output serving on Trn3. The implementation includes features such as LNC=1 decode and a decode megakernel, among other optimizations. The megakernel fuses all decoder layers of a decode step into one NKI kernel. The implementation runs as a disaggregated deployment, with six prefill replicas and one decode server on two Trn3 instances. The Overview explains each part, and the steps run the configuration with self-contained scripts.

Known limitations:

The implementation supports Trn3 with MXFP4 weights only.
The megakernel runs only on the decode server of a disaggregated deployment. A single-instance server and the prefill servers use the per-layer path at LNC=2.
The tested shape is 8,192 input tokens and 1,024 output tokens. The servers use --max-model-len 10240, a block size of 128, and an FP8 KV cache.
The decode server uses --max-num-seqs 8 on each of its 16 data-parallel ranks. This gives 128 concurrent requests for each decode instance.

Use this tutorial for GPT-OSS 120B with MXFP4 weights on Trn3 at 8K/1K. For a different case, use the gpt-oss tutorial. It covers gpt-oss 20B, Trn2 with BF16 weights, a single-instance deployment, and a context of up to 16,384 tokens.

Overview#

The configuration is a disaggregated-inference (DI) deployment with six prefill replicas and one decode server (6P1D). It runs on two Trn3 instances.

Node	Role	Servers	Parallelism	Logical NeuronCores
A	Prefill	6 replicas on ports 8100 to 8105	TP4 DP1 EP2 for each replica, at `NEURON_LOGICAL_NC_CONFIG=2`	24 of 64
A	Proxy and benchmark client	Proxy on port 8000
B	Decode	1 server on port 8200	TP8 DP16 EP128, at `NEURON_LOGICAL_NC_CONFIG=1`	128 of 128

Note

The prefill instance of this deployment has idle NeuronCores. The six replicas use 24 of its 64 logical NeuronCores. For production, keep six prefill replicas for each decode instance, and fill each prefill instance with 16 replicas. Three prefill instances and eight decode instances use every NeuronCore. That deployment has 48 prefill replicas and 8 decode servers on 11 instances.

The decode server differs from the GPT-OSS DI deployment in the gpt-oss tutorial in these points:

Part	What it does
LNC=1 decode	`NEURON_LOGICAL_NC_CONFIG=1` gives one logical NeuronCore for each physical core. The decode server runs 128 ranks, with one expert on each rank (EP128).
Decode megakernel	One NKI kernel runs all decoder layers of a decode step. vLLM Neuron selects it automatically for an MXFP4 decode server at LNC=1 on Trn3. No flag turns it on.
Packed FP8 KV cache	`--kv-cache-dtype fp8` with `fp8_packed_kv` on both roles, and a block size of 128.
Fused token embedding	The megakernel builds the token embedding (`embedding_dp_size: 1` and `VLLM_NEURON_MEGA_FUSE_EMBED=1`).
Fused sampling tail	The final RMSNorm, the lm_head, the top-k and the sampler run in one chunked kernel, with one 128-rank all-to-all (`gpt_oss_tail_hops: [128]`) and one row gather (`VLLM_NEURON_TAIL_FLAT_ROW_GATHER=1`).
DMA-order schedule	A schedule file sets the order of the DMA transfers inside the megakernel (`VLLM_NEURON_MEGA_DMA_ORDER_CONSTRAINTS_FILE`).
Host-side settings	Four batches in flight (`VLLM_NEURON_BATCH_QUEUE_DEPTH=4`), a backoff on the completion poll, and a delay on the NIXL progress thread.

The scripts of this tutorial are in the vLLM Neuron repository, in examples/vllm_neuron/models/gpt_oss/120b/mxfp4/multinode_DI_8k1k/. They launch each part with the tested configuration. The directory is self-contained: it carries the DMA-order schedule and the helper that prepares it.

File	Where to run	Purpose
`env_common.sh`	Sourced by every script	The addresses, the shape, the ports, and the environment that both roles need
`run_decode.sh`	Node B	The megakernel decode server (`kv_consumer`) with its tuning
`run_prefill.sh <index>`	Node A	One prefill replica (`kv_producer`), index 0 to 5
`run_prefill_fleet.sh`	Node A	All six replicas, one after the other
`run_proxy.sh`	Node A	The example proxy
`run_client.sh`	Node A	A health check and one completion through the proxy
`run_bench.sh`	Node A	The two benchmark legs: a warmup leg and the measured leg
`prepare_megakernel_schedule.py`	Called by `run_decode.sh`	Writes the schedule for the kernel-call scope of the active compile route
`gpt_oss_120b_megakernel_dma_order.json`	Read by the helper	The DMA-order schedule of the decode megakernel for Neuron SDK 2.33

Before you start#

This tutorial assumes that you have worked through:

The gpt-oss deployment tutorial. It explains the GPT-OSS DI deployment and its flags.
The disaggregated inference tutorial. It explains the NIXL KV connector, the proxy, and the EFA setup.

Tested versions#

Component	Version
Neuron SDK	2.33
vLLM Neuron plugin	Shipped with the Neuron SDK 2.33 release
vLLM (upstream)	0.29.0
gpt-oss 120B checkpoint	`openai/gpt-oss-120b` (Hugging Face)

Prerequisites#

Two Trn3 instances with EFA connectivity. Put them in the same subnet and placement group. The security group must allow this traffic between the two instances: port 8000 (proxy), ports 8100 to 8105 (prefill replicas), port 8200 (decode), the NIXL side-channel ports 5559, 5659, 5759, 5859, 5959 and 6059 (prefill) and 8759 (decode), and all EFA traffic (a self-referencing rule).
Neuron SDK 2.33 with vLLM Neuron on both instances. See the setup guide. Install nixl with its LIBFABRIC runtime libraries, as the DI tutorial describes.
The vLLM Neuron examples on both instances. Clone the vLLM Neuron repository, or copy its examples/vllm_neuron/ directory. The scripts are in examples/vllm_neuron/models/gpt_oss/120b/mxfp4/multinode_DI_8k1k/, and the proxy is in examples/vllm_neuron/vllm/disaggregated_inference/.
The checkpoint on a local NVMe volume of both instances. Each server loads the checkpoint at startup. A local NVMe volume gives a faster startup than a network volume or a download at startup. Download openai/gpt-oss-120b from Hugging Face to the same path on both instances, and set MODEL_ID to that path in step 1. Put the compile cache on local NVMe too, as the setup guide describes.
```
hf download openai/gpt-oss-120b --local-dir /path/to/nvme/gpt-oss-120b
```

Step 1: Set the addresses#

In this step, you tell the scripts where the two nodes are.

On both nodes, export the private IP of each node:

export PREFILL_HOST=10.0.1.10                # node A
export DECODE_HOST=10.0.1.20                 # node B
export MODEL_ID=/path/to/nvme/gpt-oss-120b   # or openai/gpt-oss-120b

Replace the two addresses with the private IPv4 addresses of your instances. Every script sources env_common.sh. That file reads these variables and exports the environment that both roles need. Every other value has a default. Open the file to see each value and the reason for it.

The logs of the servers go to ~/gpt-oss-120b-8k1k-results/logs/, and the benchmark results go to ~/gpt-oss-120b-8k1k-results/. Set BENCH_RESULT_DIR or SERVER_LOG_DIR to change the location.

Step 2: Launch the decode server on node B#

In this step, you start the one decode server. It takes every NeuronCore of the node.

cd examples/vllm_neuron/models/gpt_oss/120b/mxfp4/multinode_DI_8k1k
./run_decode.sh

The script does three things before it calls vllm serve:

It sets NEURON_LOGICAL_NC_CONFIG=1 and the decode-only environment.
It writes the DMA-order schedule for the decode graph into MEGAKERNEL_SCHEDULE_DIR (default ~/.cache/vllm-neuron/megakernel_schedule), and it exports the path of the file. The path is part of the compile-cache key, so keep one directory.
It exports VLLM_NEURON_MEGA_FUSE_EMBED=1 and VLLM_NEURON_TAIL_FLAT_ROW_GATHER=1.

The first launch compiles the decode graph. Later launches use the compile cache. The server reports Application startup complete when it is ready. You can start the prefill replicas (step 3) at the same time.

Step 3: Launch the six prefill replicas on node A#

In this step, you start six TP4 prefill servers. Replica i binds NeuronCores 4i to 4i+3, serves HTTP on port 8100 + i, and uses the NIXL side-channel port 5559 + 100 * i.

cd examples/vllm_neuron/models/gpt_oss/120b/mxfp4/multinode_DI_8k1k
./run_prefill_fleet.sh

The fleet script starts replica 0 and waits until it answers /health, so the compile cache holds the prefill graph. Then it starts replicas 1 to 5 the same way. Replica 0 compiles the prefill graph, so it takes the longest to start. The later replicas use the compile cache. The replicas keep running after the script exits. Their logs are in ~/gpt-oss-120b-8k1k-results/logs/prefill_<i>.log. To start one replica by hand, run ./run_prefill.sh <index>.

The six replicas are six separate TP4 DP1 EP2 servers, and not one DP6 server. Expert parallelism cannot share data-parallel ranks on the prefill graph.

Step 4: Launch the proxy on node A#

In this step, you start the proxy and send one request through it. Start the proxy after the decode server and all six replicas report Application startup complete:

./run_proxy.sh &
./run_client.sh

The health check of the proxy (/healthcheck) returns HTTP/1.1 200 OK, and the completion returns an OpenAI-compatible JSON payload with a coherent choices[0].text.

Note

The proxy is toy_proxy_server.py from the vLLM Neuron examples. It sends the prompts to the prefill replicas in turn (round robin). It is a minimal example router, and not a production router. For production, use a production router.

Step 5: Run the benchmark#

In this step, you run two benchmark legs against the proxy. Both legs send random prompts of 8,192 tokens, ask for 1,024 output tokens with --ignore-eos, and keep 128 requests in flight. 128 is the decode capacity: 8 sequences on each of the 16 data-parallel ranks.

./run_bench.sh

Leg	Prompts	Purpose
Warmup	1,280	The first leg after a server launch is not a valid measurement. Discard it.
Measured	12,800	The result. The script prints the key metrics and keeps `bench_clean.json` in `BENCH_RESULT_DIR`.

Before the measured leg sends its first request, the client generates the 12,800 random prompts of 8,192 tokens. That step takes several minutes, and the servers are idle while it runs. Read these keys of bench_clean.json:

median_tpot_ms and p99_tpot_ms: the time for each output token, as the client sees it. 1000 / median_tpot_ms is the output tokens per second for each user.
output_throughput: the output tokens per second over all 128 users.
total_token_throughput: the input plus output tokens per second.
mean_ttft_ms and p99_ttft_ms: the time to the first token. The prefill replicas and the KV transfer set it.

Configuration reference#

Shape#

These values are the tested configuration of the megakernel decode.

Setting	Value	Reason
`--max-model-len`	10240	Holds the 8,192 input tokens and the 1,024 output tokens.
`--max-num-batched-tokens`	8192	One prefill segment holds the whole prompt. Both roles use one bucket, `[8192]`, for the KV segment and for the batched tokens.
`--block-size`	128	The KV block size of the packed FP8 cache.
`--kv-cache-dtype fp8` and `fp8_packed_kv: true`	Both roles	The packed FP8 KV layout that the prefill writes and that the megakernel reads.
`--max-num-seqs` (decode)	8	The partition axis caps the per-rank decode batch at `128 * TP / EP = 8`. A larger value stops the trace with `Partition dim size must be preserved`.
`--max-num-seqs` (prefill)	1	One prompt for each prefill step.
`--data-parallel-size 16`, `--enable-expert-parallel`, `ep_degree: 128` (decode)		16 TP8 ranks give 128 ranks, with one expert on each rank.
`embedding_dp_size: 1` (decode)		The megakernel builds the token embedding, so the embedding is not sharded over the DP group.
`lm_head_dp_size: 16` (decode)		The lm_head is replicated over the 16 data-parallel ranks.
`gpt_oss_tail_hops: [128]` (decode)		One 128-rank all-to-all in the fused tail.
`--optimization-level 2` (decode)		The compiler optimization level of the decode graph.
`--max-logprobs 0`	Both roles	No logprob gather.
`--no-disable-hybrid-kv-cache-manager`	Both roles	Keeps the hybrid KV cache manager that the sliding-window layers need.
`--hf-overrides '{"quantization_config": {}}'`	Both roles	Lets vLLM load the published MXFP4 checkpoint. `quantization: "mxfp4"` in `neuron_config` selects the Neuron MXFP4 path.
`--distributed-executor-backend mp`	Both roles	The multiprocess executor.

Environment#

Both roles, from env_common.sh:

Variable	Value	Reason
`VLLM_NEURON_BACKEND`	`neuron_native`	The compile route. libtorch-neuronx-lite reads this variable itself, and it selects its native-route setup only for an explicit value.
`NEURON_LIBTORCH_DISABLE_GRAPH_CAPTURE_BACKEND`	`1`	The native route needs the graph-capture backend off.
`TORCH_NEURONX_VOCAB_SHARDING_SPMD_DISABLE`	`1`	Turns off the vocab-sharding FX pass, so the per-rank vocab offset stays out of the graph.
`NEURON_RT_DBG_VECTOR_DGE_SKIP_NOTIF`	`1`	Packed-FP8 KV writes benign out-of-bound skip notifications. Without this setting, the notifications stop the decode server on a long run.
`VLLM_NEURON_NIXL_PROGRESS_THREAD_DELAY_US`	`250`	The NIXL progress thread polls with no delay by default and takes a core from the worker.
`VLLM_NIXL_SIDE_CHANNEL_HOST`	`0.0.0.0`	The side channel must accept the KV handshake from the other node.
`NEURON_LIBTORCH_COMPILATION_TIMEOUT`	`2400`	The limit in seconds for the compile of one graph. The default is 600, and the compile of one prefill graph takes several minutes.
`VLLM_ENGINE_READY_TIMEOUT_S`	`14400`	The limit in seconds for the startup of one server. The default is 600, and the first prefill replica compiles its graph during the startup.

Decode only, from run_decode.sh:

Variable	Value	Reason
`NEURON_LOGICAL_NC_CONFIG`	`1`	The megakernel needs one logical NeuronCore for each physical core.
`VLLM_NEURON_MEGA_COLLECTIVE_PRIVATE_SRC`	`1`	Keeps the source tensors of a collective off shared HBM. A toolchain that requires this stops the trace without it.
`VLLM_NEURON_BATCH_QUEUE_DEPTH`	`4`	Four batches in flight hide the host time of a step.
`NEURON_RT_XU_WORKER_ENABLE_LOW_LATENCY_COMPLETION`	`0`	The completion poll backs off and gives host cores back to the workers.
`VLLM_NEURON_MEGA_DMA_ORDER_CONSTRAINTS_FILE`	The schedule path	The DMA-order schedule. When this variable is set, vLLM Neuron also sets the kernel options that the schedule was made with.
`VLLM_NEURON_MEGA_FUSE_EMBED`	`1`	The token embedding inside the megakernel.
`VLLM_NEURON_TAIL_FLAT_ROW_GATHER`	`1`	One row gather in the fused tail.
`VLLM_NIXL_SIDE_CHANNEL_PORT`	`8759`	Past the last prefill side-channel port.

Prefill only, from run_prefill.sh:

Variable	Value	Reason
`NEURON_LOGICAL_NC_CONFIG`	`2`	The prefill graph runs the per-layer path at LNC=2.
`NEURON_VISIBLE_DEVICES`	`4i-(4i+3)`	The four cores of replica `i`.
`VLLM_NIXL_SIDE_CHANNEL_PORT`	`5559 + 100 * i`	One side-channel port for each replica. The same value goes into `nixl_side_channel_port` of the additional config.

Script options:

Variable	Default	Purpose
`MODEL_ID`	`openai/gpt-oss-120b`	A Hugging Face model ID, or the local path of the checkpoint.
`BENCH_RESULT_DIR`	`~/gpt-oss-120b-8k1k-results`	The benchmark results.
`SERVER_LOG_DIR`	`$BENCH_RESULT_DIR/logs`	The logs of the servers and of the proxy.
`MEGAKERNEL_SCHEDULE`	The JSON file of this directory	The shipped DMA-order schedule.
`MEGAKERNEL_SCHEDULE_DIR`	`~/.cache/vllm-neuron/megakernel_schedule`	Where `run_decode.sh` writes the schedule copy. Keep one directory: its path is in the compile-cache key.
`MEGAKERNEL_SCOPE`	`auto`	The kernel-call scope of the schedule. `auto` gives 250 on the native route and 261 on the XLA route, the values of the Neuron SDK 2.33 release.
`PROXY_SCRIPT`	The proxy of the examples tree	The path of `toy_proxy_server.py`, if you copied the directory somewhere else.
`PROXY_BIND_HOST`	`0.0.0.0`	The address that the proxy listens on.
`PROXY_URL`	`http://localhost:8000`	The proxy address for `run_client.sh` and `run_bench.sh`.
`BENCH_NETRC`	`/dev/null`	The netrc file that the benchmark client reads before each request. Keep a local path. See Common issues.

Common issues#

The benchmark reports a low throughput and a high TTFT, but the servers are healthy#

The benchmark client can be the bottleneck. vllm bench serve opens its HTTP session with trust_env, so the client looks up ~/.netrc before it sends each request. On a network file system, each lookup can take hundreds of milliseconds. The client then sends requests late. The request throughput decreases and the TTFT increases, but the TPOT and the server metrics stay normal.

run_bench.sh points the lookup at /dev/null (BENCH_NETRC). If you run vllm bench serve yourself, set NETRC=/dev/null on the command. You can also run the client as a user whose $HOME is on a local disk.

The decode TPOT is much larger than expected#

Start the decode server with run_decode.sh. The script sets the megakernel tuning. Each setting is off by default, and without them the server gives correct output at a larger step time.

After a change to a kernel option or to the schedule file, clear the compile caches on both nodes. The cache key does not include the content of the schedule file or the kernel options.

rm -rf ~/.cache/neuron_libtorch ~/.cache/vllm/neuron /var/tmp/neuron-compile-cache /var/tmp/nki-*cache*

The decode trace stops with `Partition dim size must be preserved`#

--max-num-seqs on the decode server is above 8. Keep 8 for TP8 and EP128.

The decode server stops with `At LNC=1 a collective's src tensors must not be on shared_hbm`#

VLLM_NEURON_MEGA_COLLECTIVE_PRIVATE_SRC=1 did not reach the server. Run the server through run_decode.sh.

A server stops during compilation#

A compile or a startup took longer than its limit. Raise NEURON_LIBTORCH_COMPILATION_TIMEOUT, the limit for one graph, and VLLM_ENGINE_READY_TIMEOUT_S, the limit for the startup of one server, in env_common.sh.

Requests stall after the prefill, or the NIXL handshake fails#

Check VLLM_NIXL_SIDE_CHANNEL_HOST=0.0.0.0 on both roles. Check the side-channel ports in the security group: 5559 to 6059 in steps of 100 for the replicas, and 8759 for the decode server. Check the EFA rule. Enable the NIXL logs as the disaggregated inference tutorial describes.

The decode server stops on a long run with many out-of-bound skip notifications#

NEURON_RT_DBG_VECTOR_DGE_SKIP_NOTIF=1 is absent on one role. env_common.sh sets it for both roles.

Clean up#

Stop each server with SIGINT, so it shuts down in order:

pkill -INT -f "vllm serve"       # on both nodes
pkill -f toy_proxy_server        # on node A