From d49180019bb25cbb7ef2daf533d972a5e86d737e Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 5 Sep 2026 17:13:07 +0800 Subject: [PATCH] fix(moe): cast filtered-activation expert_ids to int32 for torch.compile (#38085) Co-authored-by: BBuf --- .../kernels/ops/activation/activation.py | 5 ++++ .../kernels/ops/activation/test_activation.py | 30 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/python/sglang/kernels/ops/activation/activation.py b/python/sglang/kernels/ops/activation/activation.py index f5765181e..41e63514c 100644 --- a/python/sglang/kernels/ops/activation/activation.py +++ b/python/sglang/kernels/ops/activation/activation.py @@ -134,6 +134,11 @@ def run_activation( if expert_ids is None: _run_activation_inplace(op_name, input, out) else: + # The JIT kernel indexes expert ids as int32. Routing ids may arrive as + # int64 (e.g. from torch.topk) and torch.compile realizes them at their + # true dtype, so normalize here instead of asserting downstream. + if expert_ids.dtype != torch.int32: + expert_ids = expert_ids.to(torch.int32) _run_activation_filtered_inplace(op_name, input, out, expert_ids, expert_step) return out diff --git a/test/registered/kernels/ops/activation/test_activation.py b/test/registered/kernels/ops/activation/test_activation.py index a191cdbdb..fb2562e0f 100644 --- a/test/registered/kernels/ops/activation/test_activation.py +++ b/test/registered/kernels/ops/activation/test_activation.py @@ -166,6 +166,36 @@ def test_activation_filter_expert_none_skipped(op_name: str) -> None: torch.testing.assert_close(out_filtered, out_unfiltered, atol=0.0, rtol=0.0) +def test_activation_filter_expert_int64_under_torch_compile() -> None: + """torch.topk routing ids remain usable after AOTAutograd realizes int64.""" + shape = (32, 512) + dtype = torch.bfloat16 + x = torch.randn(shape, dtype=dtype, device="cuda") + expert_ids = torch.zeros((shape[0],), dtype=torch.int64, device="cuda") + expert_ids[::3] = -1 + out = torch.full( + shape[:-1] + (shape[-1] // 2,), + float("nan"), + dtype=dtype, + device="cuda", + ) + + def compiled_activation(input, output, routing_ids): + return run_activation("silu", input, output, routing_ids, 1) + + result = torch.compile(compiled_activation, fullgraph=True)(x, out, expert_ids) + assert result is out + + skipped = expert_ids == -1 + assert torch.isnan(out[skipped]).all() + torch.testing.assert_close( + out[~skipped], + _reference("silu", x)[~skipped], + atol=1e-2, + rtol=1e-2, + ) + + UNARY_SHAPES = get_ci_test_range( full_range=[ (7, 16),