From 283d76d56a45d9bbead5ba650402fb0957abc30c Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Tue, 6 Oct 2026 21:45:04 -0400 Subject: [PATCH] [Bugfix] Widen ep_gather output offset to int64 (#58373) Signed-off-by: Robert Shaw Co-authored-by: Claude Code --- tests/kernels/moe/test_deepgemm.py | 22 +++++++++++++++++++ .../layers/fused_moe/deep_gemm_utils.py | 2 +- 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 385cc596ed..e341f12a71 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -70,6 +70,28 @@ def test_ep_gather_uses_64_bit_row_offsets(): torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0) +def test_ep_gather_uses_64_bit_output_row_offsets(): + hidden_size = 6144 + output_row = (1 << 31) // hidden_size + 1 + num_rows = output_row + 1 + output = torch.empty( + (num_rows, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + + ep_gather( + input_tensor=torch.ones((1, hidden_size), device="cuda", dtype=torch.bfloat16), + recv_topk_ids=torch.zeros((num_rows, 1), device="cuda", dtype=torch.int64), + recv_topk_weight=torch.ones((num_rows, 1), device="cuda", dtype=torch.float32), + input_index=torch.zeros((num_rows, 1), device="cuda", dtype=torch.int32), + expert_map=None, + output_tensor=output, + ) + + torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0) + + @pytest.mark.parametrize("workspace_dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("num_columns", [2048, 6144, 6145]) def test_fp8_workspace_shape(workspace_dtype, num_columns): diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index 624aa9ce86..1dee9bad54 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -415,7 +415,7 @@ def _fwd_kernel_ep_gather( tl.store( output_tensor - + cur_token * output_tensor_stride0 + + tl.cast(cur_token, tl.int64) * output_tensor_stride0 + cur_block * BLOCK_D + off_d, accumulator.to(output_tensor.dtype.element_ty),