feat(examples): add gpt-oss-120b example - #83
Open
ymwangg wants to merge 4 commits into
Open
Conversation
ymwangg
force-pushed
the
feat/gpt-oss-120b-example
branch
from
August 24, 2026 18:16
f28661e to
332fc58
Compare
A NKIPy implementation of OpenAI's gpt-oss-120b MoE model for AWS Trainium (TRN2), with tensor/expert/data parallelism (TP/EP/DP). - Standalone nki (beta-3) kernels throughout: rmsnorm, router, fused rank-slice, fused rmsnorm-gemm, blockwise MoE, and the flash-attention prefill/decode stack; no neuronxcc dependency. - HuggingFace gpt_oss checkpoint support as the default weight-prep path (MXFP4 dequant + reshard), with the OpenAI block.* checkpoint path kept as an alternative. Co-authored-by: Zhenyu Song <zhenyus@amazon.com> Co-authored-by: Lingfan Yu <lingfany@amazon.com> Co-authored-by: Yanming Wang <yanmwang@amazon.com> Co-authored-by: Ziyang Xu <ziyangx@amazon.com>
ymwangg
force-pushed
the
feat/gpt-oss-120b-example
branch
from
August 27, 2026 21:36
332fc58 to
c50b3fa
Compare
Async decode (opt-in via GPT_OSS_120B_ASYNC=1): submit the per-step decode kernel chain (embedding -> fused decode blocks -> sampling) non-blocking on the runtime FIFO channel with a bounded in-flight window (GPT_OSS_120B_ASYNC_WINDOW, default 4), waiting on the oldest in-flight future before the window overflows and on all outstanding futures (in FIFO order) at step end. Overlaps host dispatch with device compute; ~1.30x decode at batch=1 / 64-rank (33.5 -> 43.7 tok/s, decode MBU 19.1% -> 24.9%), output byte-identical. Uses DeviceKernel.submit / SpikeAsyncFuture.wait from the runtime (#84), the same non-blocking submit/wait pattern the gpt-oss-20b async decode path uses. Launcher / usability fixes: - chat.py: tolerate --checkpoint given with or without the -TP{tp} suffix (and a trailing slash), with a real error message instead of a bare AssertionError. - chat.sh: activate the repo venv (so python AND neuronx-cc resolve there), default to the single-node-verified layout (DP8/EP8, was DP16/EP4), and guard the port-cleanup lsof so it doesn't abort under set -e. - config.py: drop NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS (removed from NRT; only produced a per-rank warning). - README: document the layout defaults, the 128-rank (LNC1 + ep8) recipe, and the HBM/LNC sizing notes.
ymwangg
force-pushed
the
feat/gpt-oss-120b-example
branch
from
August 27, 2026 23:23
c50b3fa to
4ae36d6
Compare
blockwise_nki_static fully unrolled its block loop (Python range over num_static_blocks, ~159 in the 64-rank layout), which dominated compile: recompiled once per prefill layer, it was ~91% of "Preparing kernels". When BUFFER_DEGREE == 1 (the prefill path) there is no cross-block buffer rotation to unroll for, so iterate blocks with a real device loop (nl.dynamic_range) instead. block_idx becomes a runtime nisa.VirtualRegister, so the two block-indexed gathers (load_token_indices, _load_block_expert) now select the block via a DMA scalar_offset rather than a static element offset. BUFFER_DEGREE > 1 (tokengen/decode) keeps the static unroll so the compiler can rotate buffer_idx = block_idx % BUFFER_DEGREE and software-pipeline weight DMA; that kernel (blockwise_nki_tokengen_one_tile_replicated_hidden_state) is untouched. Measured (trn2, 64-rank DP8/EP8/TP8): "Preparing kernels" ~15 min -> ~1.8 min (~8x). Prefill costs ~15% (TTFT 0.39 -> 0.46s at a 960-token prompt) from the lost cross-block overlap; decode is a different kernel and unchanged. Unit tests (test_blockwise_nki, bf16 + float8, prefill + tokengen) pass; full-model generation is byte-identical to the unrolled kernel.
Flip GPT_OSS_120B_ASYNC to default on (was opt-in). The decode step already submits its kernel chain (embedding -> fused decode blocks -> sampling) non-blocking on the runtime FIFO channel via DeviceKernel.submit / SpikeAsyncFuture.wait, overlapping host dispatch with device compute; making it the default gives that speedup out of the box. Set GPT_OSS_120B_ASYNC=0 to fall back to the blocking path for reference / debugging. Measured (trn2, 64-rank DP8/EP8/TP8, batch=1): decode 33.8 -> 44.1 tok/s (~1.30x), byte-identical output. Verified as the default (no env vars set) end to end, together with the device-loop prefill kernel.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
A NKIPy implementation of OpenAI's gpt-oss-120b MoE model for AWS Trainium (TRN2), as a new
examples/models/gpt_oss_120bexample separate from the existing 20b example. Unlike the 20b example (dense/batched numpy MoE, TP-only), this targets scale: blockwise MoE (NKI kernel + C++ pybind11 index builder), TP + prefill expert-parallel + data-parallel, NKI prefill/decode attention, and a 4-layer fused decode kernel.Scope is core model only (no vLLM plugin). Verified end-to-end on trn2.48xlarge (coherent generation — "The capital of France is Paris") with a presharded TP8 checkpoint.
Highlights
nki(beta-3) kernels throughout — rmsnorm, router, fused rank-slice, fused rmsnorm-gemm, blockwise MoE, and the flash-attention prefill/decode stack. Noneuronxccdependency remains. All device-validated; decode throughput recovered to baseline after migration (single-DMA sin/cos rope gather).scripts/hf_tensor_preparation.pyreads a HuggingFacegpt_osscheckpoint (e.g.openai/gpt-oss-120b), dequantizes the MXFP4 MoE weights to bf16, remaps to the internal layout, and shards for TP. The HF→internal mapping is validated byte-identical to the original OpenAIblock.*pipeline (openai_tensor_preparation.py), which is kept as an alternative.GPT_OSS_120B_ASYNC=1) — submits the per-step decode kernel chain (embedding → fused decode blocks → sampling) non-blocking on the runtime FIFO channel viaDeviceKernel.submit/SpikeAsyncFuture.wait(feat(runtime): non-blocking submit for DeviceKernel via SpikeAsync #84), with a bounded in-flight window (GPT_OSS_120B_ASYNC_WINDOW, default 4). This is the same non-blocking submit/wait pattern the gpt-oss-20b example uses. Overlaps host dispatch with device compute; output is byte-identical to the blocking path.GPT_OSS_CORE_STRIDE) so ranks spread across devices instead of packing onto a few (avoids HBM OOM).Performance
Measured on trn2.48xlarge, 64-rank layout (DP8 / EP8 / TP8), batch=1, single prompt, 64 generated tokens:
__call__)GPT_OSS_120B_ASYNC=1)Async decode gives a 1.30× speedup at batch=1 while producing byte-identical output. TTFT is unchanged (~0.37s).
Authors
Co-authored by: Zhenyu Song (zhenyus@amazon.com), Lingfan Yu (lingfany@amazon.com), Yanming Wang (yanmwang@amazon.com), Ziyang Xu (ziyangx@amazon.com).
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.