[WIP] feat(gpt-oss-megakernel): fused decode megakernel for GPT-OSS-120B - #988
Draft
vmalepati1 wants to merge 1 commit into
Draft
[WIP] feat(gpt-oss-megakernel): fused decode megakernel for GPT-OSS-120B#988vmalepati1 wants to merge 1 commit into
vmalepati1 wants to merge 1 commit into
Conversation
Add a persistent Triton megakernel that runs all decoder layers and the
final RMSNorm of GPT-OSS-120B in a single launch on gfx950 (MI350), for
world-size 1 decode. Gated behind --enable-gpt-oss-megakernel and applied
only when the model shape matches (64 attention heads, hidden size 2880).
When it cannot run, forward falls back to the standard per-op path.
Three env-selected variants share one weight snapshot and dispatch path:
- precise: fp32 dequant (default when the megakernel is enabled)
- MFMA (TOKENSPEED_GPT_OSS_MEGAKERNEL_MFMA=1): MFMA dot path
- leader-follower barrier (TOKENSPEED_GPT_OSS_MEGAKERNEL_MSBAR=1):
adds a two-phase grid barrier (sharded arrival atomic plus a single
published release flag), supersedes MFMA when set
The variants are bit-exact to the shipped dot path at bs 1, 4, 8 across all
positions. Measured decode TPOT on MI350X, world-size 1: precise 19.1 ms,
leader-follower barrier 9.7 ms.
Wires two server args (enable_gpt_oss_megakernel, gpt_oss_megakernel_max_bs)
through server_args and the global env dict.
Signed-off-by: vmalepati1 <malepativikas@gmail.com>
borontion
changed the title
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.
Summary
Adds a persistent Triton decode megakernel for GPT-OSS-120B that runs all decoder layers and the final RMSNorm in a single kernel launch on gfx950 (MI350), for world-size 1 decode.
It is gated behind
--enable-gpt-oss-megakerneland only applies when the model shape matches (64 attention heads, hidden size 2880). When it cannot run,forwardfalls back to the standard per-op path.Variants
Three env-selected variants share one weight snapshot and dispatch path:
TOKENSPEED_GPT_OSS_MEGAKERNEL_MFMA=1): MFMA dot pathTOKENSPEED_GPT_OSS_MEGAKERNEL_MSBAR=1): adds a two-phase grid barrier (sharded arrival atomic plus a single published release flag), supersedes MFMA when setCorrectness and performance
The variants are bit-exact to the shipped dot path at bs 1, 4, and 8 across all positions.
Measured decode TPOT on MI350X, world-size 1:
Changes
python/tokenspeed/runtime/models/:gpt_oss_megakernel/,gpt_oss_megakernel_mfma/,gpt_oss_megakernel_mfma_msbar/gpt_oss.py:forwardoverride that dispatches to the megakernel when it can run, plus runner selection at weight-load timeenable_gpt_oss_megakernel,gpt_oss_megakernel_max_bs) wired throughserver_argsand the global env dictNotes
Draft. Tests and docs are not yet included.