feat(moe): select gfx950 TP A8W4 SiTU

[panditsa]

Summary

• register A8W4 SiTU only for the exact Kimi K3 TP8/EP1 shard shape
• forward optional shared-down tensors through the MoE apply API
• require exact ispp traits when a kernel declares one

Performance

4K/1K TP8/EP1 CUDA graphs, median of three at head e689319e:

M	Output tok/s	TTFT	TPOT
2	90.60	659.66 ms	21.451 ms
4	173.36	1139.20 ms	21.977 ms

The exact parent does not support this Kimi SiTU TP8/EP1 configuration: it falls into the generic weight-only MXFP4 path and raises Weight-only MXFP4 MoE supports only SwiGLU. This PR therefore enables the configuration; there is no valid parent E2E percentage comparison. The direct kernel comparison is documented in #1138.

Tests

• TP SiTU reference numerics, explicit output reuse, shared-down fusion, and CUDA graphs: 7 passed
• exact-ispp selector coverage: 1 passed
• all 18 measured 4K/1K requests completed with valid 4096-token prompts and 1024-token outputs

Stack

• depends on feat(amd): add gfx950 A8W4 SiTU core #1138

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bcebe08a43

ℹ️ 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".

[chatgpt-codex-connector]

Document the new shared-projection API

Document these parameters in the exported tokenspeed_kernel.moe_apply docstring, including their all-or-none requirement, expected shapes, backend limitations, and the tuple return produced when shared projection is enabled. Without that contract, callers cannot safely use the newly public arguments or determine when the return type changes from a tensor to (routed, shared).

AGENTS.md reference: AGENTS.md:L45-L47

Useful? React with 👍 / 👎.

[nirvedhmeshram]

Hi you might already be aware but, just FYI, using the tp8ep1 (A8W4) I got accuracy issues

#	prompt	expected	--ep-size 8 (A16W4)	--ep-size 1 (A8W4)
1	What is 17 * 24?	408	408	no answer — 512 tok, finish_reason=length, repetition loop
2	Compute 2^16.	65536	65536	echoes the prompt back
3	Capital of Australia?	Canberra	Canberra	no answer — 512 tok, finish_reason=length
4	First 8 prime numbers	2, 3, 5, 7, 11, 13, 17, 19	2, 3, 5, 7, 11, 13, 17, 19	2, 3, 5, 7, 9, 11, 13, 17
5	120 km in 1.5 h, speed?	80	80	no answer — 512 tok, finish_reason=length
6	One sentence: why is the sky blue?	(one sentence)	correct	no answer — 512 tok, finish_reason=length
7	Capital of France?	Paris	Paris	echoes the prompt back
8	Is 91 prime?	no, 7 x 13	No... 91 = 7 x 13	no answer — 512 tok, finish_reason=length

EP8: 8/8 correct. EP1: 3/8 terminate at all, and all three are wrong — 0/8 correct.

Here is my experimental branch based on this work, that captures the state as it was yesterday, if it has changed since then, then please ignore this
https://github.com/nirvedhmeshram/tokenspeed/tree/a8w4-tp8ep1

[panditsa]

Hi you might already be aware but, just FYI, using the tp8ep1 (A8W4) I got accuracy issues

prompt expected --ep-size 8 (A16W4) --ep-size 1 (A8W4)

1 What is 17 * 24? 408 408 no answer — 512 tok, finish_reason=length, repetition loop
2 Compute 2^16. 65536 65536 echoes the prompt back
3 Capital of Australia? Canberra Canberra no answer — 512 tok, finish_reason=length
4 First 8 prime numbers 2, 3, 5, 7, 11, 13, 17, 19 2, 3, 5, 7, 11, 13, 17, 19 2, 3, 5, 7, 9, 11, 13, 17
5 120 km in 1.5 h, speed? 80 80 no answer — 512 tok, finish_reason=length
6 One sentence: why is the sky blue? (one sentence) correct no answer — 512 tok, finish_reason=length
7 Capital of France? Paris Paris echoes the prompt back
8 Is 91 prime? no, 7 x 13 No... 91 = 7 x 13 no answer — 512 tok, finish_reason=length

EP8: 8/8 correct. EP1: 3/8 terminate at all, and all three are wrong — 0/8 correct.

Here is my experimental branch based on this work, that captures the state as it was yesterday, if it has changed since then, then please ignore this https://github.com/nirvedhmeshram/tokenspeed/tree/a8w4-tp8ep1

This issue is now fixed.