[HIP] [CK] [MoE] Reject non-int32 index buffers in the topk kernels

[i-kosarev]

Problem

topk_softmax, grouped_topk and biased_grouped_topk write their outputs through reinterpret_cast<int*>(...) without checking the tensor's dtype:

// csrc/kernels/topk_softmax_kernels.cu
reinterpret_cast<int*>(topk_indices.data_ptr()),

A caller that allocates a 64-bit index buffer gets only the low half of every slot written; the rest keeps whatever the allocation happened to contain. Nothing fails — the call returns normally with a plausible-looking tensor.

Measurement

gfx950 (MI350X), one direct topk_softmax call, buffer pre-filled with -999999 so an untouched slot is unmistakable. 256 tokens, topk 6, 64 experts:

buffer dtype	ids landing in [0, 64)	min / max
torch.int32	1536 / 1536 (100%)	0 / 63
torch.int64	14 / 1536 (0.9%)	-999999 / 270582939709

The surviving poison is the proof those slots were never written.

repro

import torch
from aiter import topk_softmax

M, E, TOPK = 256, 64, 6
gating = torch.randn(M, E, device="cuda", dtype=torch.float32)

for dt in (torch.int32, torch.int64):
    w = torch.empty(M, TOPK, device="cuda", dtype=torch.float32)
    ids = torch.empty(M, TOPK, device="cuda", dtype=dt).fill_(-999999)
    tei = torch.empty(M, TOPK, device="cuda", dtype=torch.int32)
    topk_softmax(w, ids, tei, gating, False, 0, "")
    torch.cuda.synchronize()
    ok = int(((ids >= 0) & (ids < E)).sum())
    print(f"{dt}: in_range={ok}/{ids.numel()} min={int(ids.min())} max={int(ids.max())}")

Why it matters

Downstream this corrupts whatever consumes the ids, far from the kernel that produced them. In vLLM's ROCm expert-parallel path it reached DeepEP, which indexes with them, producing two unrelated-looking crash signatures from one cause: an out-of-bounds write (Memory access fault ... Write access to a read-only page), and — once the obviously invalid ids were clamped — duplicate expert ids per token tripping DeepEP's own device assert and aborting with HSA_STATUS_ERROR_EXCEPTION. Tracking that back to a dtype mismatch in the routing kernel took a while.

The change

Add the dtype checks that topk_gating already performs for exactly these arguments:

AITER_CHECK(topk_weights.dtype() == AITER_DTYPE_fp32, "topk_weights must be float32");
AITER_CHECK(topk_indices.dtype() == AITER_DTYPE_i32,  "topk_indices must be int32");

So this brings its siblings in line rather than introducing a new contract.

Compatibility

No behaviour change for correct callers — every allocation in this repository already uses int32:

• aiter/fused_moe.py:4261,4279
• aiter/fused_moe_bf16_asm.py:608
• op_tests/test_topk_softmax.py:22-23, op_tests/test_moe_topk_gating.py

The caller-side fix for the vLLM case is vllm-project/vllm#55147. This change is what turns the same mistake into an immediate, attributable error for the next caller instead of silent memory corruption.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

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

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5255 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[Copilot]

🟡 Changes recommended

The new fail-fast behavior is not covered by negative tests (e.g., asserting int64 index buffers raise), which risks regressions back to silent corruption.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds runtime dtype validation to MoE top-k routing kernels to prevent silent memory corruption when callers provide non-int32 index output buffers, aligning behavior with existing topk_gating contract checks.

Changes:

• Add AITER_CHECK dtype guards to topk_softmax for topk_weights (fp32), topk_indices (int32), and token_expert_indices (int32).
• Add AITER_CHECK dtype guards to grouped_topk / biased_grouped_topk for topk_weights (fp32) and topk_ids (int32).

File summaries

File	Description
csrc/kernels/topk_softmax_kernels.cu	Adds fp32/int32 dtype checks before writing outputs via reinterpret_cast<int*>.
csrc/kernels/topk_softmax_kernels_group.cu	Adds fp32/int32 dtype checks for grouped routing variants that write topk_ids via reinterpret_cast<int*>.

Review details

Suppressed comments (1)

csrc/kernels/topk_softmax_kernels_group.cu:1297

• Similarly for grouped_topk, there’s no test coverage asserting that passing a non-int32 topk_ids buffer fails fast. Adding an explicit negative test would help ensure this ABI constraint doesn’t accidentally get relaxed or broken later (especially since the kernel writes through reinterpret_cast<int*>).

    const aiter_tensor_t& correction_bias = topk_ids;
    AITER_CHECK(topk_weights.dtype() == AITER_DTYPE_fp32,
                "topk_weights must be float32");
    AITER_CHECK(topk_ids.dtype() == AITER_DTYPE_i32,
                "topk_ids must be int32");

• Files reviewed: 2/2 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟡 Changes recommended

Related ASM, sigmoid, and fused-gate routes remain unguarded, and the rejection behavior lacks tests.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

Suppressed comments (2)

csrc/kernels/topk_softmax_kernels.cu:804

• The existing op_tests/test_topk_softmax.py coverage only allocates the three outputs with their valid dtypes, so none of these new rejection contracts is exercised. Add negative tests using pytest.raises(RuntimeError, match=...) for non-float32 weights and non-int32 values of each index buffer; otherwise these guards can regress while the current correctness matrix stays green.

    AITER_CHECK(topk_weights.dtype() == AITER_DTYPE_fp32,
                "topk_weights must be float32");
    AITER_CHECK(topk_indices.dtype() == AITER_DTYPE_i32,
                "topk_indices must be int32");
    AITER_CHECK(token_expert_indices.dtype() == AITER_DTYPE_i32,
                "token_expert_indices must be int32");

csrc/kernels/topk_softmax_kernels_group.cu:1243

• The grouped-top-k tests in op_tests/test_moeTopkSoftmax.py:293-307,403-416 only pass float32 weights and int32 IDs, so the newly added failure behavior for both grouped entry points is untested. Add explicit invalid-dtype tests, including the large-token public biased_grouped_topk fallback route, to verify every dispatch path rejects before launch.

    AITER_CHECK(topk_weights.dtype() == AITER_DTYPE_fp32,
                "topk_weights must be float32");
    AITER_CHECK(topk_ids.dtype() == AITER_DTYPE_i32,
                "topk_ids must be int32");

• Files reviewed: 2/2 changed files
• Comments generated: 3
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The new test file lacks an entry point, so the standard CI invocation executes none of its tests.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

The new sigmoid fp32 restriction rejects fp16 and bf16 weight outputs that its existing CK trait dispatch explicitly supports.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

csrc/py_itfs_ck/topk_sigmoid_kernels.cu:25

• This rejects output dtypes that this CK wrapper explicitly supports: lines 43-66 map topk_weights.dtype() across fp16/bf16/fp32 and pass it as weight_prec so the kernel writes with the selected width. Consequently, existing topk_sigmoid callers using fp16 or bf16 weights now fail even though only the index buffer has the fixed 32-bit contract. Keep the int32 index check, but remove this fp32-only restriction and update the comment accordingly.

• Files reviewed: 6/6 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

The guards span multiple GPU backends and the updated branch lacks complete head-matched runtime validation.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Copilot]

🟡 Changes recommended

The ASM path unnecessarily rejects non-int32 token_expert_indices despite never consuming that buffer.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 5/5 changed files
• Comments generated: 1
• Review effort level: Balanced

[Copilot]

🔵 Needs a closer look

The current tree lacks regression tests exercising the newly added rejection paths.

Review details

Suppressed comments (1)

csrc/kernels/topk_softmax_kernels.cu:804

• The promised regression test is absent from the current tree: op_tests/test_topk_index_dtype_contract.py no longer exists, and the existing top-k tests only allocate int32 outputs. Consequently none of the new rejection paths—including the ASM, CK sigmoid, grouped, and fused-gate routes—is exercised, so these guards could regress while the positive tests remain green. Restore the subprocess-based negative cases for every changed entry point (including the fused-gate redispatch route).

    AITER_CHECK(topk_indices.dtype() == AITER_DTYPE_i32,
                "topk_indices must be int32");
    AITER_CHECK(token_expert_indices.dtype() == AITER_DTYPE_i32,
                "token_expert_indices must be int32");

• Files reviewed: 5/5 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[i-kosarev]

@valarLip @zufayu @junhaha666 any chance you could also look at #5256 ?