fix: align ck_moe_stage1 split-K tmp_out buffer with CK kernel

[karverma-amd]

Motivation

The splitK optimization in aiter's CK MoE kernel reduces average kernel time from 44us to 22us (2x improvement). However, enabling it in aiter commit: f80d9c0 causes:

RuntimeError: HIP runtime error: invalid argument. hip_check_error.hpp: 18
in function: hip_check_error

Technical Details

The crash occurs during CUDA graph capture in vLLM when ksplit > 1 for decode batches with small token counts (token_num=1, topk=8).

Root Cause

Buffer overflow in tmp_out allocation in ck_moe_stage1() (/app/aiter/aiter/fused_moe.py line ~1638).

The Mismatch

When splitK is active (KBatch > 1), the CK C++ kernel calls hipMemsetAsync to zero-initialize the output buffer for atomic accumulation. The size it zeroes is:

M * N * sizeof(float) * 2

where M = sorted_size = min(token_num * topk * block_m, sorted_token_ids.shape[0]).

But Python allocated tmp_out as:

torch.zeros((token_num, topk, w1.shape[1]), dtype=fp32)

For DeepSeek V3 decode with token_num=1, topk=8, block_m=16:

• Python buffer: 1 * 8 * 4096 * 4 = 131,072 bytes (128 KB)
• C++ hipMemsetAsync: 128 * 2048 * 4 * 2 = 2,097,152 bytes (2 MB)
• Overflow: 16x beyond the allocated tensor

C++ Code Reference

The overflow originates in device_moe_gemm_blockscale.hpp (CK):

if(arg.KBatch > 1)
    hipGetErrorString(hipMemsetAsync(arg.p_c_grid, 0,
        arg.M * arg.N * sizeof(CDataType) * (IsInputGemm && IsSplitK ? 2 : 1),
        stream_config.stream_id_));

Fix

Replace the tmp_out allocation to match the kernel's expected buffer size:

# BEFORE (buggy):
tmp_out = torch.zeros(
    (token_num, topk, w1.shape[1]), dtype=dtypes.fp32, device=out.device
)

# AFTER (fixed):
sorted_size = min(token_num * topk * block_m, sorted_token_ids.shape[0])
tmp_out = torch.zeros(
    (sorted_size, w1.shape[1]), dtype=dtypes.fp32, device=out.device
)

Test Result

DeepSeek-R1-0528 with aiter commit - f80d9c00b9c8 runtime error is resolved.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[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 2508 --add-label <label>

[ChuanLi1101]

Looks good with minor suggestions. Pre Approved for speed up the deliveries.

[ChuanLi1101]

cktile_moe_stage1 has the identical undersized buffer pattern for split_k > 1. It calls aiter.moe_cktile2stages_gemm1 which is a different kernel, but if that kernel has the same hipMemsetAsync behavior, it would have the same overflow. The TODO at line 1723 (# TODO: support fp32 splitk) suggests this path may not be actively used yet, but it's worth either applying the same fix or adding a guard/comment.

[valarLip]

not good to me, looks like we do zeros twice, 1 here, 1 ck, we need remove one

[rbrugaro-amd]

This ROCm/rocm-libraries#5225 already addressed the proper allocation from the C++ side. I have swap to torch.empty to avoid the double zeroing #2551 (created new PR because i cannot edit current branch)
@ChuanLi1101 i see you also created another PR but I think the double zeroing still there?

[karverma-amd]

Alternative PR merged - #2551