Skip to content

feat(amd): prepare gfx950 SiTU prefill - #1137

Open
panditsa wants to merge 6 commits into
lightseekorg:mainfrom
panditsa:sanket/wip-k3-situ-foundation
Open

feat(amd): prepare gfx950 SiTU prefill#1137
panditsa wants to merge 6 commits into
lightseekorg:mainfrom
panditsa:sanket/wip-k3-situ-foundation

Conversation

@panditsa

@panditsa panditsa commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • add a SiTU-v2 epilogue option to the gfx950 package-prefill stage
  • preserve logical W13 K metadata while retaining the padded W2 aliases required by stage 2
  • keep both aliases as zero-copy views over the preshuffled weights

Performance scope

This is a foundation-only stack entry. At this boundary no production caller passes situ_linear_beta, so DO_SITU remains false and there is no standalone end-to-end performance claim. #1138 adds the A8W4 core and #1139 registers/selects it.

Tests

  • gfx950 MXFP4 regression suite: 4 passed, 6 skipped
  • metadata check confirms W13 exposes logical K, W2 retains physical padded K, and both aliases preserve the original data pointers
  • downstream selector and end-to-end coverage is in (WIP) feat(moe): select gfx950 TP A8W4 SiTU #1139

Stack

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Project shared weights across M=1/2/4 while scheduling per-token AttnRes reductions, with an MFMA projection at M=4.

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Use fused Gluon for M=1/2 and the composed MFMA path from M=3 through M=16.

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
Add the SiTU-v2 epilogue to the package prefill stage and preserve logical K metadata across padded weight aliases.

Signed-off-by: Sanket Pandit <sanket.pandit@amd.com>
@panditsa
panditsa force-pushed the sanket/wip-k3-situ-foundation branch from 3785c40 to 2c32487 Compare August 19, 2026 20:21
@panditsa
panditsa force-pushed the sanket/wip-k3-latent-projection branch from 6406405 to 3a351c3 Compare August 19, 2026 20:24
@panditsa
panditsa changed the base branch from sanket/wip-k3-latent-projection to main August 25, 2026 01:59
@panditsa panditsa changed the title (WIP) feat(amd): prepare gfx950 SiTU prefill feat(amd): prepare gfx950 SiTU prefill Aug 25, 2026
@panditsa
panditsa marked this pull request as ready for review August 25, 2026 02:00
@panditsa
panditsa requested review from a team as code owners August 25, 2026 02:00

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 58b5342883

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment on lines +65 to +68
k_packed = (
int(getattr(weight, "original_k_pk", k_packed_padded))
if preserve_logical_k
else k_packed_padded

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep non-tile-aligned W13 shapes off this kernel

When the logical packed K is not a multiple of 128 (for example, GPT-OSS hidden size 2880 gives original_k_pk=1440), shortening this alias makes the package eligibility check in fused/moe.py pass. However, prefill_stage1.py uses cdiv(K, 256) and performs unmasked 128-byte activation and eight-scale loads for the final tile, while _quantize_mxfp4_activation allocates only the logical 1440 packed bytes and 90 scales. Batches of at least nine tokens can therefore read past the quantized activation buffers; either pad/mask the activation tail or keep these shapes on the fallback path.

Useful? React with 👍 / 👎.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant