[WIP] Feat/gpt oss example#63
Draft
ymwangg wants to merge 29 commits into
Draft
Conversation
Add a NKIPy example for OpenAI's gpt-oss MoE models (gpt-oss-20b / 120b), mirroring the qwen3 example structure. The implementation is fully config-driven, so both sizes share one codebase. gpt-oss-specific handling: - MXFP4 experts dequantized to bf16 at prep time - interleaved gate/up de-interleaved at prep time - clamped SwiGLU with gate_up/down biases - per-head attention sinks + QKV/O biases (no QK-norm) - alternating sliding-window / full attention (one kernel per type) - YaRN RoPE (inv_freq precomputed from HF config) - router with top-k-then-softmax and router bias Validated against HF on trn2 (TP=4): every generated token matches HF's argmax or a bf16-resolution tie.
Implements parallel-drafting P-EAGLE (arXiv 2602.01469) on top of the gpt-oss base model for speculative decoding on Trainium. Components added (examples/models/gpt_oss/eagle/): - config.py: EagleConfig for the 4-layer P-EAGLE drafter (llama3 RoPE, fc fusion, mask_hidden/ptd_token_id, d2t vocab map) - tensor_preparation.py: convert P-EAGLE checkpoint to x@W form (replicated) - kernels/drafter.py: parallel-drafting forward - K tokens in one pass via NTP (real hidden) + MTP (mask_hidden) positions with cross-depth mask - kernels/drafter_layer.py: EAGLE-3 fusion midlayer + plain Llama layers - kernels/verify.py: multi-position greedy argmax for verification - drafter_model.py: device-side drafter model + compile - speculate.py: full speculation loop (prefill → draft → verify → accept) Base model changes: - config.py: added aux_layers config + default_aux_layers() for EAGLE-3 taps - gpt_oss.py: run_prefill() now optionally captures pre-layer hidden states at the 3 EAGLE-3 tap layers (2, L/2, L-3) - kernels/attention.py: generalized decode path to support seq_len>1 (for the multi-token verify pass) via query_pos = start_pos + arange(seq_len) Status: functionally correct (lossless greedy output verified against HF). Acceptance length is below the paper's reported ~3.3 — under investigation (likely a hidden-state position/timing issue in the draft-verify loop seeding).
…yers Switch aux capture to post-layer (output of tap layers 2/12/21) based on HF validation showing the drafter predicts correctly with HF's hidden states at hs[3]/hs[13]/hs[22] (output of layers 2/12/21). Note: acceptance length remains low (~1.0) due to numerical divergence between nkipy's Neuron-compiled target and the HF CPU reference the drafter was trained against. The drafter kernel is mathematically correct (validated against independent torch reference) and correctly predicts the target when fed exact HF hidden states. The gap is an implementation-coupling issue inherent to EAGLE-style speculation.
Key findings from the P-EAGLE paper (Figure 2, Figure 3, Section 3):
1. The drafter maintains its own KV cache across the full context
(prompt + all accepted tokens). At each draft step, K positions
attend to the FULL accumulated cache.
2. The attention mask is GROUP-CAUSAL: all K positions see the full
cache (group 0), but within the K positions the NTP (group 1)
and MTP (group 2+) positions use cross-depth causality — MTP
positions cannot attend to positions at the same or later depth.
3. The NTP pair is (emb(t_n), hidden_after_processing_t_{n-1}),
predicting t_{n+1}. The hidden is one step behind the embedding.
This commit adds:
- drafter_cpu.py: CPU reference drafter with full KV cache and
standard causal attention (working infrastructure, mask needs
the group-causal refinement for MTP positions)
- Fixes hidden state capture to post-layer (output of tap layers)
- Adds peagle_aux_layers config method
Status: KV cache infrastructure correct, still needs the group-causal
mask refinement for the MTP positions within the K-wide draft window.
Root-caused the low acceptance length (~1.4 vs the card's 3.30-3.80 at K=7) on GPU by running the identical checkpoint through vLLM's eagle3 parallel-drafting path, capturing its drafter I/O, and reproducing it with a standalone PyTorch reference (cosine 0.9999, 100% draft-token match). Three bugs plus a prompt-formatting issue: 1. Context-blind drafting (dominant): speculate.py drove DrafterModel (kernels/drafter.py), which runs only the K draft positions under a (K,K) cross-depth mask with no prefill and no KV cache, so the MTP slots never saw the prompt. Rewired speculate.py to use the KV-cached DrafterCPU: prefill the drafter on the prompt (EAGLE +1 shift), then each step roll the cache back to the last accepted position and run [newly-accepted tokens | K-1 ptd slots] in one parallel forward attending to the full context. 2. rollback() truncated the wrong axis: the cache is (B, n_kv, seq, head_dim) and rollback sliced dim 1 (n_kv) instead of dim 2 (seq), so rejected speculative KV was never discarded and corrupted later steps. 3. Aux tap off-by-one: vLLM's eagle3 default (2, n//2, n-3) captures the residual stream entering those layers; our post-layer capture must shift down one, so default_aux_layers now returns (1, 11, 20). Verified on GPU the drafter's 3 fc chunks equal target layer outputs (1, 11, 20) at cos 1.0. 4. Prompt formatting: the drafter is trained on chat data; raw prompts roughly halve acceptance (GPU, K=7: 3.65 chat vs 1.99 raw). speculate.py now applies the chat template by default (--raw-prompt to opt out). Still produces all K draft tokens in a single forward pass (parallel drafting). Adds test_drafter_cpu.py guarding the rollback/full-context invariants (skips without the checkpoint). Validated against vLLM on GPU; not yet re-validated on Trainium, and the on-device kernels/drafter.py KV-cache port remains follow-up. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
Author
gpt-oss + P-EAGLE — Implementation Status ReportStatus as of the current branch TL;DR
Components
What changed recently (this work)The drafter had a working CPU reference but the on-device path was the old
Port design note: the device drafter keeps static kernel shapes — it commits Also fixed along the way: Validation
Under greedy verification emitted tokens are the target's argmax regardless of Performance (5-prompt sweep, n=128, K=7, TP=4, chat template, identical env)
(Full per-prompt table in Known issues / open work
|
The device drafter was context-blind (no KV cache, no prefill): its MTP slots never saw the prompt, collapsing acceptance to ~1.5 vs the CPU reference's ~4.3. Port the CPU drafter's full-context KV-cache algorithm to device, mirroring the base model's on-device cache pattern. - drafter_layer.py: add _attention_cached + drafter_layer_cached (KV-cached attention, causal over absolute positions). - kernels/drafter.py: add drafter_layer_kernel (per-layer cached device entry) + drafter_head_kernel. - drafter_model.py: rewrite DrafterModel with per-layer KV caches and prefill()/draft() matching DrafterCPU; implicit rollback via greedy overwrite + causal mask. - speculate.py: device drafter is now the default; --cpu-drafter selects the PyTorch reference. Also fix apply_chat_template crash on transformers 5.5.4 (return_dict=True). - run_drafter_device.py: standalone harness; multi-step logit-level parity check vs CPU. Validated on trn2 (TP=4, K=7): device matches CPU logits at cosine 0.9999 (token diffs are sub-noise bf16 ties); acceptance 4.19 vs 4.33, decode 34.8 vs 4.8 tok/s (~7x faster, no per-step host<->device sync).
Profiling (SPEC_PROFILE=1) showed the device drafter spent ~30% of decode
time in the commit step: each accepted token was pushed through all 4
drafter layers as a separate width-1 kernel launch (~4 launches/step),
paying fixed dispatch overhead for near-zero compute.
Replace the commit loop + separate draft pass with ONE combined forward of
width W = C+K-1 over [commit_0..commit_{C-1} | ptd_0..ptd_{K-2}], matching
DrafterCPU.draft. Compile one kernel set per width (W in K..2K) upfront;
lazy in-loop compilation briefly regressed decode to 0.55 tok/s.
Measured on trn2 (TP=4, K=7): 34.5 -> 41.7 tok/s, acceptance 4.19 -> 4.33
(the single pass is more faithful than commit-then-draft). Drafter time
52 -> 37 ms/step; verify (67 ms) is now the dominant cost. Parity with CPU
preserved (logit cosine 0.9999). Profiling instrumentation is gated behind
SPEC_PROFILE=1 (zero overhead when unset).
Profiling (SPEC_PROFILE=1) showed the drafter's device compute was only ~9ms/step, but the per-layer host loop (4 kernel launches + separate head + host fc-fuse + host concat + 201k-vocab logit transfer) added ~24ms of overhead, making the drafter cost as much per layer as the 24-layer target. Replace the host-driven per-layer loop with drafter_fused_kernel: one device call doing fc-fusion + [embed|hidden] assembly + all 4 layers (each aliasing its own KV cache) + final norm + lm_head. DrafterModel now uploads inputs once and invokes one kernel per step; fc-fuse moved on-device. Measured on trn2 (TP=4, K=7): 40 -> 50 tok/s, drafter forward 33 -> 17 ms/step, acceptance held at 4.33. Parity with CPU preserved (logit cosine 0.9999; token diffs are sub-noise bf16 ties). Verify (66ms) is now the dominant cost at ~78%; profiling confirms it is genuine device layer compute, not the host MoE loop the README previously (wrongly) blamed. Also remove the now-dead pre-fusion kernels (drafter_kernel/forward, drafter_layer_kernel, drafter_head_kernel, non-cached drafter_layer/ _attention) and correct the README notes.
Debugged the verify step, which is the dominant per-step cost. Measured its scaling by sweeping K on trn2 (TP=4): verify layer time is ~19 ms fixed + ~6.3 ms/candidate token. The fixed ~19 ms equals a single-token forward (HBM weight streaming), amortized across the K+1 batch — so verify is efficient per token and NOT dominated by per-token MoE weight re-gather as previously assumed. The real lever is K vs the acceptance plateau. This drafter saturates around 3.3 accepted tokens/step, so K>3 buys verify time for no extra accepted tokens: K=1 1.94acc 50.0 t/s | K=2 2.67 52.6 | K=3 3.31 ~60 | K=4 3.31 49.5 | K=7 4.33 50.4 K=3 is the throughput optimum at ~60-62 t/s (n=256), beating both the base model (~52) and the old K=7 default (~50). Change default to K=3 and document the sweep + scaling model in the README.
Contributor
Author
Progress update: P-EAGLE speculative decoding — working + optimized on trn2Brought the on-device P-EAGLE drafter from broken → 62 tok/s (vs 52 tok/s for the base model), and profiled the verify step to a firm conclusion. All changes are on this branch (4 commits). What changed
Parity held throughout: device drafter matches the CPU reference at logit cosine 0.9999 (token diffs are sub-noise bf16 ties). CPU drafter tests 3/3. Throughput progression (trn2, TP=4, K=7, chat prompt)
Latency breakdown (per verify step, SPEC_PROFILE=1)At the fused/optimized state, K=7 (~116 ms/step): Reference points:
Choosing K — why 3, not 7Verify cost grows ~6.3 ms/token, but this drafter's acceptance plateaus at ~3.3, so K>3 buys verify time for no extra accepted tokens. Sweep (n=96):
Can verify be sped up with a better kernel? Investigated → no (for this framework)Verify is 78% of each step, so I profiled the MoE hard:
Conclusion: verify is near-optimal on this framework — no worthwhile static MoE-kernel win. Grouped MoE would help large-L prefill/TTFT ( Remaining levers (algorithmic, not kernel)
Profiling is behind |
Contributor
Author
Follow-up: verify latency breakdown — where the ~63 ms goesAttributed the verify layer's cost to its components by microbenchmarking each at the real gpt-oss-20b sharded dims (TP=4, S=K+1=8 tokens), × 24 layers:
(Isolated components sum higher than the fused 63 ms because fusion overlaps DMA with compute and shares intermediates — the proportions are the point.) Which kernel to focus on: the MoE expert FFN loop, and only that. Attention (~10%) and the router (~8%) aren't worth touching. Important caveatThe per-expert matmul is not slow — it's already efficient. The cost is structural: the verify batch does
So, to keep improving verify
Net: verify is near-optimal for what this static-trace framework can express today; the remaining wins are framework or drafter improvements, not a rewrite of the expert kernel. |
Rename examples/models/gpt_oss/eagle/ to peagle/ to reflect the parallel-EAGLE (P-EAGLE) drafter it implements, updating the eagle.* package imports and path references accordingly. Algorithm-name mentions (eagle3, get_eagle_config) are left unchanged.
Adds moe_batched() to kernels/feedforward.py: expresses all N=B*L tokens x top_k experts as a single gather (np.take) + batched matmul instead of the per-(token,expert) Python loop. Numerically identical to the baseline (max|delta|=0 in bf16). Includes _moe_bench.py (isolated device-kernel A/B + equivalence check) and _ntff_stats.py (device exec-time extraction from .ntff traces). Finding: only ~1.09-1.13x on device across L=1/4/8. neuronx-cc lowers the batched dot back into one distinct-weight GEMV per (token,expert) -- matmul count still scales linearly with N*top_k -- so there is no PE-level fusion. The ~10% gain is only from collapsing router/topk/reduce glue and slightly less DMA. A real speedup needs a grouped/block-diagonal GEMM kernel (custom NKI), which numpy-tracing cannot express. Not yet wired into transformer_layer.
The --dtype bf16 path silently ran in fp32: `bf16_array * 0.1` promotes the array back to float32 (python-float scalar -> float64 wins the numpy promotion), undoing the .astype(DT). Scale before the final cast, assert the input dtype matches, and print it so this can't recur. Honest bf16 numbers: batched MoE is 1.24x on device at L=4 (was a 1.13x fp32 artifact), correct within bf16 tol (rel 6.6e-3). Tracing the bf16 kernel also shows weight DMA = 199.7MB = the 16 selected experts read exactly once, i.e. neuronx-cc already fuses the gather into an indexed weight load -- the only remaining headroom is GEMV->GEMM density (grouped matmul), not the gather.
Explores the GEMV-density headroom flagged in the batched MoE. Two new
variants in feedforward.py:
moe_concat - fuse a token's top_k experts into 2 matmuls (stack gate/up
on the output dim, down on the contraction dim so the
expert-sum falls out). Measures identical to batched: the
compiler already fuses the batched GEMV the same way, so
the explicit concat only adds a transpose that nets out.
moe_dense_masked - run ALL 32 experts as one dense batched GEMM, mask by the
router. ~flat ~1185us regardless of N (weight-load bound).
_moe_bench.py reworked to sweep {baseline, batched, concat, dense} x N and
report device exec time + matmul count per variant.
Result (bf16, TP4-rank dims, device speedup vs baseline loop):
N=1: batched/concat 1.25x, dense 0.23x
N=4: batched/concat 1.23x, dense 0.79x
N=8: batched/concat 1.22x, dense 1.55x
Dense crosses over batched at ~N=6. So: batched for decode / small verify
(K<=4), dense for large verify (K>=5). P-EAGLE default K=3 -> batched.
All variants numerically match baseline.
transformer_layer reads GPT_OSS_MOE_KERNEL (loop|batched|concat|dense) to pick the MoE impl; loop is the default reference. peagle/_sweep.py runs speculate.py across K x kernel, persists each combo to jsonl (resumable after a kill), and tabulates end-to-end tok/s. Sweep (TP4, n=160, raw prompt) tok/s: K loop batched dense 1 45.4 49.5 32.7 3 44.9 47.8 42.1 5 36.9 39.2 42.3 7 33.1 35.3 40.0 batched beats loop at every K (~1.07-1.09x e2e); dense crosses over at K=5 (verify N=6) as the microbench predicted and wins clearly at K=7 (40.0 vs 35.3). Best overall on this prompt: K=1 + batched, 49.5 tok/s. Practical rule: small K -> batched, K>=5 -> dense. Acceptance (hence absolute tok/s) is prompt/chat- template dependent; the per-K kernel ranking holds regardless.
Add a "Speedup over the base model" section: P-EAGLE K=3 is 1.18x (loop) / 1.27x (batched MoE) over the ~52 tok/s base, measured on the binary-search chat prompt at acceptance ~3.2. Explain why the speedup is far below the acceptance length (verify runs K+1 tokens through the full target; decode is dispatch- bound) and that speedup scales ~linearly with acceptance, which is prompt- dependent. Add a "MoE kernel selection" section documenting GPT_OSS_MOE_KERNEL (loop|batched|dense|concat), the end-to-end tok/s-by-K table, and the rule: batched for K<=4, dense for K>=5.
The base model also honors GPT_OSS_MOE_KERNEL, so dividing spec-batched (66.2) by base-loop (~52) overstated the win. Measured base decode with each kernel: loop 52.0, batched 54.9. Kernel-matched speedup is 1.18x (loop) / 1.21x (batched) -- batched lifts both base and spec, so it barely moves the ratio.
Drop the MoE microbenchmark and its ntff-stats helper; scrub the dangling references in peagle/README.md and transformer_layer.py. The end-to-end peagle/_sweep.py remains for measuring kernel choices.
Merge-readiness cleanup for the gpt-oss example:
- Replace dev-box-specific /home/ubuntu/models/... defaults with HF ids across
speculate.py, _sweep.py, run_drafter_device.py, test_drafter_cpu.py:
--model -> openai/gpt-oss-20b, --draft-model -> amazon/GPT-OSS-20B-P-EAGLE.
- Drop moe_concat: the sweep showed it measures identical to moe_batched (the
compiler already fuses the batched GEMV), so it was dead weight. Selector is
now {batched, dense}; default stays loop.
- Complete the peagle README Files table (drafter_cpu, run_drafter_device,
test_drafter_cpu, _sweep) and remove the concat row from the MoE-kernel table.
- Remove a stray 26MB ntff.json trace dump from the working tree.
Benchmarked double-buffered vs baseline decode at batch=1: 52.2 vs 51.8 tok/s (within noise). The optimization targets per-token host round-trip overhead, but decode is layer-bound (~18ms/step) and the async runtime already overlaps the token readback, so the fused-embedding + ping-pong-buffer machinery bought nothing while adding real complexity. Collapse generate() to the single (former baseline) path and drop the now-dead _generate_double_buffered, the fused greedy_sampling_with_embedding kernel, the on-device embedding table (tok_embedding_device), the two *_embed compiled kernels, and the --no-double-buffering flag. Verified generation is unchanged (greedy "Paris", 52.3 tok/s).
Extract the hand-written NKI RMSNorm (nki_rmsnorm_kernel + stream_shuffle_ broadcast helper) out of kernels/sampling.py into a dedicated kernels/nki/ package, so the low-level NKI code is separated from the numpy-traced kernels. sampling.py now imports it from .nki.rmsnorm. Behavior unchanged (default USE_NKI_RMSNORM=True); verified generation still produces "Paris" at ~53 tok/s. Documented the new subfolder in the README Files table.
…loop->reference Extract the inline per-(token,expert) Python-loop MoE out of transformer_layer into moe_reference() in feedforward.py, so all three MoE implementations are peer functions selected uniformly through _MOE_KERNELS (no more if/else special- casing the default). transformer_layer now just dispatches the dict entry. Rename the kernel from "loop" to "reference" (env value, dict key, docstrings, sweep default, README tables) -- it names the role (the readable correctness baseline) rather than implying it's simply unsophisticated, and it's the default. Drop now-unused imports (tensor_apis, softmax_kernel, feedforward_kernel) from transformer_layer. Behavior unchanged; verified generation still yields "Paris" at ~53 tok/s on the default path.
Drop the standalone (K, MoE kernel) sweep harness, matching the earlier removal of _moe_bench.py / _ntff_stats.py -- it only reproduced the perf tables and nothing imports it. The README still documents the results and methodology, and the sweep is trivially reproducible via GPT_OSS_MOE_KERNEL + speculate.py. Scrubbed the two README references.
Remove drafter_cpu.py (DrafterCPU) and the two files that existed only to compare against it: test_drafter_cpu.py (tested DrafterCPU exclusively) and run_drafter_device.py (a device-vs-CPU logit parity harness). The on-device DrafterModel is the sole drafter now. speculate.py: drop the DrafterCPU import, the --cpu-drafter flag, and the if/else branch; the device adapter is used unconditionally. Scrub DrafterCPU mentions from drafter_model.py docstrings and the peagle README (drafter- placement section, Files table, Validation table, root-cause narrative). The GPU-vs-vLLM validation history is retained, just no longer attributed to a shipped CPU file. Verified P-EAGLE still runs: acceptance 3.20 at K=3, ~56 tok/s (unchanged).
Document what C = len(commit_token_ids) means -- the number of tokens the target accepted in the previous verify (acceptance length, 1..K+1) -- why the step commits them (they carry the target hiddens that only exist post-verify), and why the fused width is C+K-1 rather than C+K (the NTP draft reuses the last committed row, so only K-1 placeholder slots are appended).
Rename _dt -> _to_device, self.m/self.p -> fusion_weights/plain_weights, and the _prepare_tensors weights param. Remove the unused cache_len field (and the reset() that only zeroed it), the unused module-level BUILD_DIR, and fold the cache-clearing rationale into the prefill() docstring. Add a C <= K+1 assert in draft() so an out-of-range acceptance length fails at the contract instead of via a bare KeyError. Fix the class docstring to reflect that draft kernels compile at load time while prefill compiles lazily. Verified on trn2 (TP=4, K=3): 60.97 tok/s, acceptance 3.20 -- unchanged.
The RoPE cos/sin table (a compile-time constant) was recomputed and re-gathered to the query positions inside every layer -- 4 identical materialize+gather pairs per drafter forward, since all layers share the same rope params. Compute it once in drafter_fused_kernel and thread the gathered (W, head_dim/2) freqs into all layers. _attention_cached / drafter_layer_cached now take freqs_cos/freqs_sin instead of raw rope params; the dead is_prefill branch is dropped (the fused kernel always passes a real start_pos, 0 for prefill, so the decode-branch gather/mask already subsume prefill). Also document in draft() why a step re-feeds all C accepted tokens (not just the newest): the prior step wrote their KV from ptd/mask placeholders, so they must be overwritten with real embeds+hidden at their absolute positions. Verified on trn2 (TP=4, K=3, n=256): 61.71 tok/s, acceptance 3.20 -- numerically identical output, ~1% faster, no extra device memory.
The draft step's natural width W = C + K - 1 varies with the acceptance count C (1..K+1), so we compiled K+1 shape-static kernels (W = K..2K) up front -- the bulk of the drafter's cold-compile time. Instead compile ONE kernel at the max width (W = 2K, C = K+1) and pad every call up to it: pad the ids with extra ptd placeholders and pad target_hidden3 with the raw mask_hidden (which fc-fuses to the same MTP hidden the kernel builds for its own MTP rows). Padding is numerically exact, not approximate: the genuine draft rows [C-1 .. C+K-2] get identical inputs and attention context, and the extra rows land at absolute positions past this step's draft window, so the causal mask hides them from every real row and their cache writes are overwritten by a later commit -- the same invariant that already lets speculative MTP rows be rewritten. Verified on trn2 (TP=4, K=3, n=256): 62.28 tok/s, acceptance 3.24 -- unchanged vs per-width. Drafter draft-kernel cold compile 180.8s -> 44.0s (~4x).
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.
Add p-eagle gpt-oss-20b example.