MegaKernels in Iris 🔥 - #541
Draft
neoblizz wants to merge 63 commits into
Draft
Conversation
Single persistent Triton kernel running attention and MoE for all 36 layers of GPT-OSS-120B on one GPU (batch-1 decode), collapsing cosmic's multi-GPU design. FP4 expert weights dequantized in-kernel, BF16 compute. Validated against a PyTorch reference: greedy decode produces coherent output on real weights. Includes HF->.iris converter and kernel tests.
Expert GEMVs can now run native FP4xFP8 scaled matrix multiply via tl.dot_scaled (lowers to v_mfma_scale_f32_16x16x128_f8f6f4 on gfx950) with dynamic FP8-E4M3 activation quant, selected by --quant. Keeps the BF16 dequant path as default. Quantized path is ~2.8x faster (29ms vs 84ms TPOT on MI355X) and still decodes coherently. Adds dot_scaled and quantized-expert correctness tests plus a TPOT benchmark.
Block-of-rows GEMV tiling with max_contiguous/multiple_of hints so the compiler emits dwordx4 loads instead of per-element ushort. Fuse the attention RMSNorm into QKV and the MoE RMSNorm into the router GEMV, fold RoPE into attention, and stripe the residual/zeroing across all programs. This removes the serial single-program phases and their grid barriers. TPOT drops from 29.4 to 25.3 ms (quant) on MI355X; output unchanged.
In the quantized path, each program now owns whole 32-element SwiGLU output blocks and quantizes the block it just produced, so the producer and consumer are the same program and the separate act-quant barrier is gone (4 fewer grid barriers per layer). TPOT drops from 25.3 to 22.9 ms (quant) on MI355X; output unchanged.
Document the measured findings (zero spills, per-element loads, barrier cost) and the tiling + barrier-reduction passes. Refresh the TPOT table: quantized 29.4 -> 22.9 ms on MI355X.
Simplify the README to match the other examples, rewrite the module and test docstrings for clarity, and remove implementation-history detail. No functional changes; both decode paths still produce the same output.
The top-k experts are independent until the final accumulation, so run each expert phase (gate-up, SwiGLU, down) across all experts before the next barrier instead of one expert at a time. This drops the expert barriers from three per expert to three per layer. TPOT 22.7 -> 19.3 ms.
The scaled FP4 x FP8 expert GEMV used the weight as the dot operand whose contiguous dimension was not the inner one, and small block sizes left most programs idle. Make the weight the lhs so its contiguous K bytes coalesce, and use BLOCK_N=32, BLOCK_K=512, which raises the expert GEMV from ~0.6 to ~1.3 TB/s. TPOT 19.1 -> 11.7 ms.
BLOCK_M=16 lifts the BF16 tiled GEMV toward 5 TB/s on the LM-head shape.
Keep the attention output in its own buffer and defer the residual add to the layer's final accumulation, so the router and expert input read rmsnorm(x + o) directly. Removes the separate residual phase and its grid barrier, and the final residual is striped across all programs. TPOT 11.8 -> 11.5 ms.
Track a running (max, index) while computing the LM-head logits instead of writing all vocabulary logits to HBM and reading them back for the argmax. Removes a full logits round-trip and one barrier. TPOT 11.5 -> 11.2 ms.
A 1024-wide K block on the scaled FP4 x FP8 expert GEMV raises its bandwidth from ~1.3 to ~1.8 TB/s. TPOT 11.2 -> 10.6 ms.
The down projection (N=H) reaches best bandwidth at BLOCK_N=16 while the gate-up (N=2*I) prefers 32, so give them independent tile sizes.
Replace the per-position scalar KV scan with a blocked online-softmax over BLOCK_T positions at a time. This removes the linear per-token slowdown as the context grows, so decode TPOT stays flat. At a 100-token context TPOT drops from ~14 to ~9.5 ms; full-rollout TPOT is ~9.4 ms.
The kernel is bound by the grid-wide barriers between phases, not by GEMV occupancy. Dropping from 256 to 192 programs makes each barrier cheaper while the per-phase GEMVs stay well filled. TPOT 9.4 -> 7.4 ms.
Compute the current position's attention term in-register from the freshly projected k/v and read only the history from the cache, so the cache append and the attention can share a phase. Removes one grid barrier per layer. TPOT 7.4 -> 7.25 ms.
Run the router GEMV, the FP8 expert-input quantization and the accumulator reset in one phase (all depend only on x + o), then compute the top-k redundantly per program so the experts need no separate top-k barrier. TPOT 7.25 -> 7.15 ms.
The last expert's down projection owns the same MoE-output rows it just wrote, so it finalizes x += o + moe in its epilogue instead of a separate striped residual phase. Removes one more grid barrier per layer (quant path). TPOT 7.15 -> 6.97 ms.
Smaller output-row tiles fill the grid better for the attention-path and LM-head GEMVs. TPOT 6.97 -> 6.78 ms.
With the reduced barrier count and BLOCK_M=8, 184 programs measures fastest. TPOT 6.78 -> 6.56 ms.
With BLOCK_M=8, a 1024-wide K block markedly speeds the attention-path and LM-head BF16 GEMVs. TPOT 6.56 -> 5.87 ms.
The LM head spans the full vocabulary, so a separate BLOCK_M_LM lets it use a larger row tile than the small attention GEMVs.
bench_islosl.py sweeps ISL:OSL pairs and reports TTFT, TPOT and end-to-end latency. README now carries a measured table (100:100, 1024:100, 1024:1024, 2048:2048) on MI355X: TPOT is flat at ~5.9 ms/token; TTFT scales linearly with input length since prefill reuses the single-token decode kernel.
Convert the inner K-reduction of every GEMV (expert FP4 dot_scaled, the BF16 attention/router GEMVs and the LM head) from a plain while-loop to a tl.range pipelined loop so the next block's loads overlap the current dot. This lifts the FP4 expert GEMV from ~1.8 to ~2.7 TB/s. TPOT 5.78 -> 5.20 ms.
github-actions
Bot
added
in-progress
neoblizz
marked this pull request as draft
June 24, 2026 15:57
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because no GitHub Actions runner was available. Make sure your repository has a runner available to run Copilot's review, or add a copilot-setup-steps.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds a GPT-OSS-120B “megakernel” example for Iris, including a single-GPU persistent Triton kernel implementation, reference math, multi-GPU prototype, and tooling for validation/benchmarking/accuracy evaluation.
Changes:
- Introduces the single-GPU persistent Triton megakernel + reusable
common/device-op library. - Adds reference implementation, phased Triton runner, and multiple validation/benchmark scripts.
- Adds a multi-GPU (1 attention/tail + 4 MoE ranks) prototype using Iris symmetric heap.
Reviewed changes
Copilot reviewed 32 out of 34 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| examples/33_gpt_oss_megakernel/tokenizer_util.py | Minimal tokenizer wrapper for HF snapshot tokenizer.json. |
| examples/33_gpt_oss_megakernel/test_quant_expert.py | Validates quantized expert GEMV path vs PyTorch reference. |
| examples/33_gpt_oss_megakernel/test_kernels.py | Validates phased Triton kernels vs reference math. |
| examples/33_gpt_oss_megakernel/test_dot_scaled.py | Validates tl.dot_scaled operand layout for FP4xFP8 GEMV. |
| examples/33_gpt_oss_megakernel/test_barrier.py | Tests grid-wide barrier semantics needed for persistent kernels. |
| examples/33_gpt_oss_megakernel/run_triton_phased.py | End-to-end decode using per-phase Triton kernels orchestrated by host. |
| examples/33_gpt_oss_megakernel/run_reference.py | End-to-end decode using PyTorch reference implementation. |
| examples/33_gpt_oss_megakernel/reference.py | Defines numerical ground-truth for decode forward pass. |
| examples/33_gpt_oss_megakernel/multi_gpu/run_multi_gpu.py | Multi-process multi-GPU driver using Iris symmetric heap exchanges. |
| examples/33_gpt_oss_megakernel/multi_gpu/protocol.py | Defines rank roles and exchange protocol for multi-GPU decode. |
| examples/33_gpt_oss_megakernel/multi_gpu/persistent_kernels.py | Persistent per-rank Triton kernels using device-side flag rendezvous. |
| examples/33_gpt_oss_megakernel/multi_gpu/moe_kernels.py | MoE-rank per-layer expert kernels + scatter-back helper. |
| examples/33_gpt_oss_megakernel/multi_gpu/attn_kernels.py | Attention-rank prologue/scatter/accumulate/lm_head kernels. |
| examples/33_gpt_oss_megakernel/load_hf.py | Loads HF checkpoint into the example’s weight layout + FP4 dequant helpers. |
| examples/33_gpt_oss_megakernel/kernels.py | Standalone Triton building blocks for phased execution. |
| examples/33_gpt_oss_megakernel/gpt_oss_120b_quantized_megakernel.py | Main single-GPU persistent megakernel + host driver. |
| examples/33_gpt_oss_megakernel/convert_to_iris.py | Converts HF checkpoint to .iris tensor archive for mmap/device loading. |
| examples/33_gpt_oss_megakernel/common/swiglu.py | Device ops for SwiGLU (BF16 + FP8-quant variants). |
| examples/33_gpt_oss_megakernel/common/router.py | Device op for top-k + softmax routing. |
| examples/33_gpt_oss_megakernel/common/rmsnorm.py | Device op to materialize residual+RMSNorm output for non-quant path. |
| examples/33_gpt_oss_megakernel/common/quant.py | Device op for fused residual+RMSNorm+FP8 activation quantization. |
| examples/33_gpt_oss_megakernel/common/gemv_fp8.py | FP8 weight-only GEMV device helpers (incl. RMSNorm/resid fused). |
| examples/33_gpt_oss_megakernel/common/gemv_fp4.py | MXFP4 expert GEMV device helpers (dequant path + dot_scaled path). |
| examples/33_gpt_oss_megakernel/common/gemv_bf16.py | BF16 GEMV device helpers (tiled + fused RMSNorm/resid RMSNorm). |
| examples/33_gpt_oss_megakernel/common/fp4.py | FP4 magnitude LUT helper for dequant path. |
| examples/33_gpt_oss_megakernel/common/barrier.py | Grid-wide barrier helper for persistent kernels. |
| examples/33_gpt_oss_megakernel/common/attention.py | Device ops for RoPE+KV append and per-head flash decode with sinks. |
| examples/33_gpt_oss_megakernel/common/init.py | Re-exports reusable device ops for megakernel and multi-GPU kernels. |
| examples/33_gpt_oss_megakernel/bench_tpot.py | TPOT benchmarking script for steady-state decode. |
| examples/33_gpt_oss_megakernel/bench_islosl.py | ISL/OSL sweep benchmark for TTFT/TPOT/E2E throughput. |
| examples/33_gpt_oss_megakernel/acc_eval.py | Accuracy comparison between BF16-attn and FP8-attn variants. |
| examples/33_gpt_oss_megakernel/README.md | Documentation: architecture, usage, benchmarking, and accuracy notes. |
| .gitignore | Adds tmp/ ignore entry. |
Comments suppressed due to low confidence (1)
examples/33_gpt_oss_megakernel/tokenizer_util.py:1
- Snapshot selection always returns the lexicographically-last snapshot directory, without checking that it actually contains
tokenizer.json. If the newest snapshot is partial/corrupt, tokenizer loading will fail later with a less actionable error. Consider filtering candidates to those containingtokenizer.json, and/or raising an error that includes attempted paths when none match.
Comment on lines
+56
to
+57
| NUM_WG = 180 | ||
| _NWG = tl.constexpr(NUM_WG) |
Comment on lines
+148
to
+157
| for t in range(lo, pos + 1): | ||
| kt = tl.load(kcache_ptr + t * kv_dim + kvh * DH + d).to(tl.float32) | ||
| score = tl.sum(q * kt, axis=0) * scale | ||
| m_new = tl.maximum(m, score) | ||
| alpha = tl.exp(m - m_new) | ||
| p = tl.exp(score - m_new) | ||
| l = l * alpha + p | ||
| vt = tl.load(vcache_ptr + t * kv_dim + kvh * DH + d).to(tl.float32) | ||
| acc = acc * alpha + p * vt | ||
| m = m_new |
Comment on lines
+15
to
+22
| tl.inline_asm_elementwise( | ||
| "buffer_inv sc0\n\ts_waitcnt vmcnt(0)", | ||
| "=r", | ||
| [], | ||
| dtype=tl.int32, | ||
| is_pure=False, | ||
| pack=1, | ||
| ) |
Comment on lines
+121
to
+122
| idx_path = os.path.join(snap, "model.safetensors.index.json") | ||
| weight_map = json.load(open(idx_path))["weight_map"] |
Comment on lines
+1
to
+3
| """Accuracy comparison: FP8-attention vs BF16-attention megakernel. | ||
|
|
||
| Both share the FP4 experts; this isolates the error introduced by quantizing the |
…date
The grid barrier ends with `buffer_inv sc0`, which corrupts decode output.
`buffer_inv` discards dirty lines rather than writing them back, and on MI355X
the L2 is per-XCD and shared with ~31 other workgroups. By the time one program
leaves the spin, peers have already passed the barrier and started writing the
next phase, and a cache-wide invalidate throws those writes away.
Measured on GPT-OSS-120B decode, 300 fresh-model reps per arm, one variable,
same node/script/prompt, each arm verified from the binary it ran:
invalidate corrupt runs distinct wrong outputs
(none) 0/300 0
buffer_inv sc0 15/300 1
buffer_inv sc1 293/300 21
buffer_inv sc0 sc1 300/300 6
Monotonic in invalidation strength, so a stronger fence makes it worse, not
better. `sc0 sc1`'s dominant failure is byte-identical to the `sc0` failure, so
scope sets the frequency of one failure mode rather than its identity. The
no-invalidate arm is Fisher one-sided p = 2.6e-05 against shipped and is
perf-neutral (4.65 -> 4.59 ms bf16; 4.40 -> 4.41 with --fp8-attn, both inside
run-to-run noise).
Add `_barrier_noinv` and use it from the single-GPU megakernel. `_barrier` is
left as-is: `multi_gpu/` still calls it, that path exchanges over the symmetric
heap rather than an intra-GPU L2, and it has not been tested against this
change. Correctness of the no-invalidate form relies on there being no
cross-phase reuse of shared addresses, which holds for BS=1 decode where
weights stream read-once; an all-to-all pattern fails 0/10 without an
invalidate. Both caveats are documented at the definition.
Co-Authored-By: Claude <noreply@anthropic.com>
Polling with `atomic_add(ptr, 0)` looks like an obvious thing to optimise: it is an unbounded RMW on a single cache line from every workgroup, and swapping it for a cache-bypassing load makes the barrier ~1.9x cheaper and the whole token ~1.2x faster. It also corrupts the output. Measured on GPT-OSS-120B decode with no other change: 25/25 runs non-reference with 25 distinct trajectories, against 0/300 with the RMW. The counting semantics are not the problem -- the counter is monotonic and the test is `>=`, so a stale poll can only read too small and spin longer, never exit early. The RMW is doing a second job: forcing a coherence point that the payload reads depend on. That is what makes a barrier with no cache invalidate sound in the first place, and a load that bypasses cache for the counter address only removes it for everything else. The scattered failure (25 distinct trajectories) is a genuine race, distinct from the single-trajectory corruption caused by `buffer_inv`. Co-Authored-By: Claude <noreply@anthropic.com>
The previous comment asserted that the atomic "is forcing a coherence point that the payload depends on". That is one of two candidate mechanisms and neither is confirmed: the alternative is that the RMW contention simply delays every workgroup long enough for peers' writes to drain, in which case the barrier is correct by being slow. Adding a single relaxed atomic back at the spin exit does not restore correctness, so it is not one discrete coherence event. A false explanation is worse than an admitted gap -- it tells the next reader they understand the constraint well enough to engineer around it. State the measurement, state that the mechanism is open, and give the gate: any change that speeds up this loop needs a >=300-rep output-correctness run. Also records why the obvious check does not work. An assert that the counter reached NB*NWG proves every arrival eventually happened; it says nothing about whether any program waited at each barrier, so workgroups can run ahead and the total still lands exactly on target. Co-Authored-By: Claude <noreply@anthropic.com>
The previous wording said correctness "relies on there being no cross-phase reuse of shared addresses -- true for BS=1 decode". The condition was right and the framing was misleading: this kernel has no batch dimension and cannot run anything but single-sequence decode, so "safe at BS=1" describes a regime it can never leave and points a reader at a test that does not apply. The audience is someone reusing this barrier in a different kernel. What they need to check is whether their programs re-read, across a barrier, shared addresses they already have cached -- not what batch size they run. State the measurement on both sides of that condition: 0/300 correct where reuse is absent, 0/10 with a 64-program all-to-all that re-reads the same slots every phase, which passes with an invalidating barrier. Co-Authored-By: Claude <noreply@anthropic.com>
The docstring reported 0/300 without saying where it came from. The run was on a node that was near-idle across the reps where failures appear, so the figure describes a quiet machine and says nothing about a busy one. That matters here because the failure being fixed is timing-sensitive. Neighbours on other GPUs cannot reach this GPU's L2, but they can lower clocks through the board power cap, and every timing in the barrier moves with them. This barrier's correctness has never been observed past the early quiet stretch under known load. No claim of a defect is intended: the only contrary observation is a single non-reference run out of 100 on a shared node, which is not distinguishable from the 0/300 at this n. Stating the conditions is true regardless of how that resolves, and a measurement without its conditions invites the reader to assume the widest one. Co-Authored-By: Claude <noreply@anthropic.com>
The previous note said behaviour under heavy node contention was uncharacterised, and gave the board power cap as the mechanism to worry about. That mechanism has since been measured and is absent: GPU 0 under this workload runs at 2388 MHz on an idle node and 2392 MHz with seven sibling GPUs saturated. Neighbours cannot reach this GPU's L2 and do not move its clock, so there is no established channel by which a co-tenant perturbs the race this barrier is exposed to. Removing that explanation makes the one contrary observation more likely to be real, not less, so the caveat now states the thing itself: pooled across every run of this barrier the rate is 1/475 (0.21%, 95% CI 0.01-1.17%), and a residual defect is not excluded. It remains the right choice by a wide margin. The barrier it replaces measured 15/300 (5.0%) on a quiet node -- roughly 24x higher, Fisher one-sided p = 5e-06. The honest framing is a large improvement with an uncharacterised remainder rather than a clean barrier. Co-Authored-By: Claude <noreply@anthropic.com>
Two problems with the figure as committed. It included a count from a run that was still executing, so the denominator was stale before the commit landed. And it presented the pooled rate ahead of the reasoning that makes pooling legitimate, so a reader who stops at the table takes a cross-node average at face value. Pooling across nodes assumes the rate does not depend on the node, which is the exact question the single event raised. That assumption is only defensible because the mechanism was measured and found absent -- neighbours do not move this GPU's clock and cannot reach its L2. If a channel is ever established, the figure has to be split by venue rather than pooled, and the note now says so. Co-Authored-By: Claude <noreply@anthropic.com>
The note listed two candidate mechanisms for why the RMW poll is load-bearing: that it orders other memory operations across the spin exit, or that its cost simply delays every workgroup long enough for peers' writes to drain. The second is now eliminated by measurement. Substituting a pure delay of 27x the RMW's measured per-barrier cost (s_sleep 4096, ~109 us against ~4.00 us) is 25/25 corrupt with 25 distinct trajectories. An earlier delay arm was dismissible on two grounds -- it ran at 43% of the dose and it shared a node -- and this one is neither: the dose is 27x and it ran exclusive, so a dirty result is interpretable. This narrows the caveat rather than widening it. The barrier's correctness is not an accident of how expensive the poll happens to be, so nothing here is contingent on timing. What remains unexplained is narrower: something about the repeated RMW's memory-model semantics, given that one atomic at the exit and a waitcnt drain both fail to substitute. Co-Authored-By: Claude <noreply@anthropic.com>
The note claimed this barrier is ~24x better than the one it replaces. Both numbers are real and the ratio between them is not: the numerator pools a quiet node with a shared one, the denominator is a quiet node only, and the two barriers have never been measured side by side at the same n on the same machine. That is the same defect as dividing a numerator and denominator from different runs, one level up -- different venues rather than different jobs. A paired run now in progress has the old barrier at 0/175 under synthetic load where it historically failed at roughly 10% past that point. That is 25 informative reps and not decisive, but it is the wrong direction to be leaning on an unpaired ratio while the paired measurement is still running. The decision this note supports does not rest on the ratio. The sc1 variants measured 293/300 and 300/300; no venue argument reaches numbers of that size, so removing the invalidate stands on its own regardless of where the low-rate figures settle. Co-Authored-By: Claude <noreply@anthropic.com>
The in-flight count this note flagged as provisional has completed: a full 0/300 on an exclusive node, which is the first clean run this barrier has had under conditions that were characterised rather than lucky. That takes the pooled figure to 1/700. Reporting it pooled hides the structure. Every run on a machine whose conditions we understand is clean -- 0/600 combined -- and the single non-reference output came from an organically busy shared node, which is the one venue we could not characterise. The note now splits by machine and gives an interval per venue. This is deliberately not phrased as reassurance. 100 reps on the busy node is thin, and the synthetic load we built to reproduce that condition was measured inert: it saturates sibling GPUs, and sibling GPUs do not move this GPU's clock (2388 MHz idle-node vs 2392 MHz stressed). A condition we cannot reproduce on demand leaves the rate there unknown rather than merely unmeasured. Co-Authored-By: Claude <noreply@anthropic.com>
This note described the barrier being replaced as measuring 15/300 (5.0%) "on a quiet node". Both halves of that are wrong. The node was not quiet: a co-tenant occupied it for the first ~117 of the 300 reps, and every one of the 15 failures landed after that neighbour left. And 5% is not the barrier's rate -- the identical file under a controlled synthetic load on another machine is 0/300, against 15/300 on the first, Fisher one-sided p = 2.6e-05 at full n on both sides. Both numbers are real measurements of the same code, so the rate depends on the machine through a channel we have not identified. It is not the clock (GPU 0 measures 2388 MHz idle-node vs 2392 MHz stressed) and it is not the L2, which neighbours cannot reach; what remains is host-side, such as launch cadence or per-rep build timing. The decision to remove the invalidate does not rest on either figure. The sc1 variants measured 293/300 and 300/300 and fire in every venue tested. Co-Authored-By: Claude <noreply@anthropic.com>
… need it The rate this change was justified by came from a single session. The identical barrier file has since run 0/300 under a controlled synthetic load and 0/300 on a node verified as sole occupant before the run started -- 600 further reps in two venues with no failures, Fisher one-sided p = 2.6e-05 against the original. There is no reproducible failure rate at the shipping dose and none should be quoted. The failures themselves were real: 15 byte-identical prompt-echo trajectories is not noise. They remain unexplained. The only pattern consistent with every run is that they began after a co-tenant left the machine mid-run, which is a transition rather than a state, and none of these measurements isolate it. Removing the invalidate remains correct on an argument that never depended on that number. Removal has zero measured cost, keeping the operation has no demonstrated benefit, the operation is destructive by mechanism and catastrophic at stronger scopes (293/300 and 300/300, reproducing in every venue at full n), and nothing has failed in any venue after removing it. No single clause carries that; together they hold without a reproducible rate. Co-Authored-By: Claude <noreply@anthropic.com>
…e one family The note said the 15/300 failed to reproduce across "two venues". Both of those venues are the same node family, and both failing runs are the same machine, so it is one family sampled twice against one machine sampled twice rather than four independent points. Sorting every run of the identical file by host rather than by load also removes the load story entirely: the machine that fails does so both with a co-tenant and under a synthetic load, and the machines that pass do so both under load and solo. A later run showed a failure at rep 50 with the load still applied, which also rules out the busy-then-idle transition that the earlier ordering suggested. Host is a label for whatever differs, not a mechanism. Slurm reports these nodes as identical in version, GRES, CPU count and memory, so the difference is below what the scheduler exposes. A run on another node of the failing family would separate "this machine" from "this family" and has not been done. Co-Authored-By: Claude <noreply@anthropic.com>
The failures of the barrier being replaced sort by host: they occur on one machine and have never occurred on the two others tested. Averaging a rate over hosts that do not exhibit the defect understates it on the host that does, and that average is what this note has been quoting. Restricting to that machine gives a comparison with no cross-venue or cross-run division on either side: the old barrier is 19/575 across two runs under two different load conditions, and the replacement is 0/600 across two runs paired in a single container. Fisher one-sided p = 1.1e-06. This is a better justification for the change than the rate it originally shipped with, because it is measured where the defect actually appears. The replacement's runs predate knowing which machine was interesting, so they were not designed as this comparison; they happen to constitute one. Co-Authored-By: Claude <noreply@anthropic.com>
The old barrier's second run on that host has finished at 4/300, so its total there is 19/600 rather than the 19/575 partial this note quoted. The replacement is 0/600 on the same host, giving equal denominators on both sides and Fisher one-sided p = 1.7e-06. The previous figure counted a run that was still executing, which is the same defect this file already corrected once for a different number. Co-Authored-By: Claude <noreply@anthropic.com>
The note said a run on another node of the failing family would separate "this machine" from "this family" and had not been done. It has now: thor-3, same family as the failing thor-2, is 0/300 at full length on a node verified as sole occupant before the run. So the old barrier's failures are specific to one machine. Load is fully eliminated as the variable -- that machine fails both with a co-tenant and under synthetic load, and every passing machine passes both loaded and solo. Stated with its power rather than as a clean result: 0/300 on thor-3 excludes thor-2's measured rate (P = 0.008 at 1.6%) but not the bottom of its confidence interval (P = 0.27 at 0.44%). And thor-2's own two runs differ by 4x, 15/300 against 4/300, so the host sorts whether it fails without producing a stable rate underneath. Co-Authored-By: Claude <noreply@anthropic.com>
The note gave one power figure, computed against a confidence floor from a single run of the failing host. Pooling that host's two runs instead moves the floor from 0.36% to 1.92% and takes the same clean 300 from "not excluded" (P = 0.33) to "decisive" (P = 0.003), with the two hosts' intervals disjoint. Neither number is wrong; they answer different questions. Pooling assumes the two runs on that host sample one rate, and they differ 4x -- 15/300 against 4/300 -- so the conservative statement is the unpooled one. Both are now given, along with the point that the gap between them is a choice about pooling rather than anything measured on the clean host. Co-Authored-By: Claude <noreply@anthropic.com>
Four small defects found while studying this example, none related to the barrier work. `--persistent` multi-GPU raises `KeyError: MEASURE_NOEXCH` before it reaches a kernel: run_multi_gpu.py passes the flag at line 227 and it is absent from attn_persistent_kernel's signature, so that path cannot run at all. Add the constexpr and let it gate the result rendezvous, which is what the flag is for. The output with it set is garbage by construction -- res[] holds the previous layer's values -- so the docstring says it is a profiling switch and not a correctness path. Add a run-time check that NUM_WG fits the device. The grid-wide barrier needs every program resident and this kernel gets one workgroup per CU, so a grid above the CU count hangs with no error rather than failing. The boundary tracks the CU count rather than a constant -- 80/81 on MI308X, 256/257 on MI355X, 304/305 on MI300X -- so it has to be checked against the device, not asserted. Add requirements.txt. The base images ship torch and triton but not tokenizers, safetensors or huggingface_hub, and the example does not run without them. Remove `bars_per_layer = 11`. Nothing reads it and it matches neither path: the quant path executes 7 barriers per layer and the bf16 path 8. A stale constant that looks like an invariant invites someone to compute a barrier target from it. Co-Authored-By: Claude <noreply@anthropic.com>
The note told a reader whose kernel re-reads shared addresses across a barrier
to use `_barrier` instead. That does not work, and the measurement contradicting
it was already in hand when the advice was written.
The 64-program all-to-all test was run against all four barrier variants. The
L1-only invalidate that `_barrier` uses fails it 0/10, exactly as no invalidate
does; only an invalidate reaching L2 passes.
no invalidate 0/10 FAIL
buffer_inv sc0 0/10 FAIL <- what _barrier does
buffer_inv sc1 10/10 pass
buffer_inv sc0 sc1 10/10 pass
So neither barrier in this file is correct for that access pattern, and pointing
at `_barrier` sends someone to a variant that fails the same way. The note now
gives the table and says so.
Co-Authored-By: Claude <noreply@anthropic.com>
The test used one monotonic counter for both edges of each phase. That is
lappable -- a fast program's own later arrival can satisfy a target meant to
require all P distinct programs -- so a program could overwrite its slot while a
slower one was still summing. It failed roughly 3 runs in 10, and because the
failure is a race rather than a wrong answer, the single run it did made that
invisible.
Add the second counter so reads complete before the next phase's writes, and run
the whole thing ten times, since one green run of a race test is not evidence.
Also record why the barrier is inlined instead of imported, which previously
looked like an oversight. Every program here re-reads all N slots every phase,
which is the cross-phase reuse the shipped barriers document as out of scope,
and both fail it:
no invalidate 0/10 FAIL
buffer_inv sc0 0/10 FAIL
buffer_inv sc1 10/10 pass
buffer_inv sc0 sc1 10/10 pass
The inlined spin uses sem="acquire", which lowers to buffer_inv sc1 on gfx950.
Importing either shipped barrier would fail, and that is a property of this
test's all-to-all sharing rather than a defect in them.
Co-Authored-By: Claude <noreply@anthropic.com>
These kernels use a grid-wide barrier at 20 call sites across three files and never check that NWG fits the device. Above the CU count some programs are never scheduled and the rest spin on a target that cannot be reached, so the failure is a hang with no error. NWG=120 fits the parts in use and not an 80-CU MI308X. The single-GPU path gained a run-time check for this; this path did not, because it has no test coverage here and adding the import blind is a worse trade than writing down the hazard. Moving the check into common/ or duplicating it locally are both reasonable, and either should be done by someone able to run the path. Recording one interaction explicitly: the --persistent entry point holds 13 of the 20 sites and used to raise KeyError before launching anything, so the hazard could not fire there. Fixing that KeyError made those sites reachable. That is still an improvement over a path that cannot run at all, but the note describes the code as it is now rather than as it was. Co-Authored-By: Claude <noreply@anthropic.com>
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
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.
This pull request introduces the GPT-OSS-120B Megakernel example, providing a comprehensive, high-performance single-GPU implementation of the GPT-OSS-120B model using a persistent Triton kernel. The changes include detailed documentation, benchmarking and accuracy scripts, and a reusable set of device-level Triton operations for attention, quantization, and expert routing. Together, these updates enable end-to-end quantized inference, benchmarking, and accuracy evaluation for this large mixture-of-experts model.
Major additions and improvements:
1. Documentation and Usage:
README.mddescribing the GPT-OSS-120B Megakernel architecture, model details, quantization options, benchmarking methodology, accuracy tradeoffs, and file purposes.2. Benchmarking and Evaluation Scripts:
bench_tpot.pyfor measuring steady-state decode latency (TPOT) and comparing quantized vs. BF16 inference paths.bench_islosl.pyto benchmark prefill and decode latency across various input/output sequence lengths, reporting throughput and latency metrics.acc_eval.pyto compare accuracy between FP8 and BF16 attention (with shared FP4 experts), providing detailed metrics such as top-1 agreement, top-k overlap, KL divergence, and logit cosine similarity.3. Reusable Device Operations:
common/__init__.pyto expose a suite of reusable Triton device ops for attention, quantization, GEMV, routing, and SwiGLU, supporting both single-GPU and future multi-GPU kernels.common/attention.pycontaining Triton JIT helpers for RoPE+KV cache appending and per-head flash decode with attention sinks, facilitating both current and future kernel designs.