Summary of Changes

Hello @ch-wan, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request is a work-in-progress (WIP) aimed at updating the Mixture of Experts (MoE) kernel, specifically focusing on the moe_align_block_size function and its underlying implementations. The core changes involve adjusting array indexing and buffer sizing in the Python, C++, and CUDA components to resolve potential off-by-one errors or to align with a revised indexing scheme for expert IDs, which is critical for accurate token routing and padding in MoE operations.

Highlights

• MoE Kernel Indexing Adjustments: I've made consistent modifications across the Python, C++, and CUDA kernels to adjust array indices (e.g., topk_ids[i] + 1, e - 1, left - 2) when processing expert IDs and token counts within the moe_align_block_size logic. This likely addresses off-by-one errors or aligns with a new indexing scheme.
• Buffer Size Correction: The cumsum_buffer in the Python layer (fused_moe.py) has been expanded from (num_experts + 1,) to (num_experts + 2,), indicating a need for an additional slot to correctly handle cumulative sums in the MoE alignment process.
• Parameter Alignment for Native Kernels: The num_experts parameter passed from the Python moe_align_block_size function to the underlying sgl_moe_align_block_size kernel is now num_experts + 1. This suggests a change in how the number of experts is interpreted or indexed by the native C++/CUDA implementations.
• Debugging Aid: A print(num_experts) statement was temporarily added in the Python moe_align_block_size function, likely for debugging purposes during the development of these changes.

Code Review

This pull request updates the Mixture-of-Experts (MoE) kernels for expert parallelism (EP). The changes consistently apply a +1 shift to expert IDs across Python, C++, and CUDA files, likely to switch to a 1-based indexing scheme internally. The logic appears consistent and correct within the scope of this change.

I've identified a leftover debug print statement in a Python file, which I've marked as high as it could cause issues in a production environment. Please ensure this is removed. Additionally, I've raised concerns about potential out-of-bounds access in C++ and CUDA code due to the -1 and -2 adjustments to expert IDs, which need careful review.

LGTM.