fix: ck_moe_stage1 split-K buffer overflow from padding scatter (alternative to #2508)

[ChuanLi1101]

Summary

Fix out-of-bounds buffer overflow in \ck_moe_stage1\ when splitK is enabled.

Root cause

The CK MoE kernel uses \sorted_size = min(token_num * topk * block_m, sorted_token_ids.shape[0])\ as its M dimension. The kernel launches tile-based blocks covering the entire M range and scatters results to the output buffer. The output buffer must be large enough to accommodate \sorted_size\ rows.

The original code allocated only (token_num, topk, w1.shape[1])\ = \ oken_num * topk\ rows (a 3D tensor). For the padding entries in \sorted_token_ids, the sentinel value (topk << 24 | token_num)\ decodes to scatter position \ oken_num * topk + topk, which exceeds the allocated buffer. Additionally, the kernel expects the output buffer to span at least \sorted_size\ rows to match its tile-based computation grid.

Fix

• Compute \sorted_size = min(token_num * topk * block_m, sorted_token_ids.shape[0])\ (matching the C++ wrapper logic)
• Allocate a 2D fp32 buffer of shape (sorted_size, w1.shape[1])\ instead of the undersized 3D (token_num, topk, w1.shape[1])\
• After the kernel, slice only the valid rows \ mp_out[:token_num*topk, :]\ before passing to \silu_and_mul\ / \gelu_and_mul\

Verification

Tested on MI355X (gfx950) with multiple token/topk/expert configurations:
\\

SplitK Scatter Fix Verification

[tok=1 topk=8 E=256] OK shape=torch.Size([1, 8, 256]) nan=False inf=False
[tok=2 topk=8 E=256] OK shape=torch.Size([2, 8, 256]) nan=False inf=False
[tok=4 topk=8 E=256] OK shape=torch.Size([4, 8, 256]) nan=False inf=False
[tok=16 topk=8 E=256] OK shape=torch.Size([16, 8, 256]) nan=False inf=False
[tok=1 topk=4 E=64] OK shape=torch.Size([1, 4, 256]) nan=False inf=False
[tok=3 topk=6 E=128] OK shape=torch.Size([3, 6, 256]) nan=False inf=False

Results: 6 passed, 0 failed out of 6
ALL TESTS PASSED!
\\

Comparison to PR #2508

PR #2508 uses \sorted_token_ids.shape[0]\ rows (safe but over-allocates). This PR uses \sorted_size\ (the exact M dimension the C++ wrapper computes), which is the minimal correct size. Both are valid; this PR is tighter on memory.

Test plan

• Verified on MI355X with 6 different token/topk/expert configs
• All tests pass with no NaN/Inf in output
• Non-splitK path is unchanged (tmp_out = out)

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-355	Run Triton tests on MI355 in addition to MI325
ci:sglang	SGLang integration tests
ci:atom	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
ci:vllm	vLLM benchmark
ci:all	All of the above

Add labels via the sidebar or gh pr edit 2547 --add-label <label>