Skip to content

Fix fused indices conversion for non-power-of-two topk/local experts #150

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhujian19891203 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix fused indices conversion for non-power-of-two topk/local experts #150

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 46 additions & 11 deletions aiak_megatron/megatron/core/fusions/fused_indices_converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
""" fused_indices_converter.py"""

import math

import torch
import triton
import triton.language as tl
Expand All @@ -17,7 +19,9 @@ def _indices_to_multihot_kernel(
probs_in_multihot_ptr,
position_map_ptr,
num_of_local_experts: tl.constexpr,
num_of_local_experts_next_power_of_2: tl.constexpr,
topk: tl.constexpr,
topk_next_power_of_2: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
'''
Expand Down Expand Up @@ -52,22 +56,31 @@ def _indices_to_multihot_kernel(
[0.0, 0.3, 0.4, 0.0]
]
'''
# Prepare the [0, topk) row
topk_row = tl.arange(0, topk_next_power_of_2)
topk_row = tl.where(topk_row < topk, topk_row, -1)
topk_row_mask = topk_row != -1
# Prepare the [0, num_of_local_experts) row
num_exp_row = tl.arange(0, num_of_local_experts_next_power_of_2)
num_exp_row = tl.where(num_exp_row < num_of_local_experts, num_exp_row, -1)
num_exp_row_mask = num_exp_row != -1

# Load a [1, topk] row from the indices buffer
row_idx = tl.program_id(0)
indices_row = tl.load(indices_ptr + row_idx * topk + tl.arange(0, topk))
probs_row = tl.load(probs_in_indices_ptr + row_idx * topk + tl.arange(0, topk))
indices_row = tl.load(indices_ptr + row_idx * topk + topk_row, mask=topk_row_mask)
indices_row = tl.where(topk_row_mask, indices_row, -1)
probs_row = tl.load(probs_in_indices_ptr + row_idx * topk + topk_row, mask=topk_row_mask)

# Get the position of the each index in the indices_row, which is saved for backwards
position_row = tl.arange(0, topk)
position_row = tl.where(indices_row != -1, position_row, -1)
position_row = tl.where(indices_row != -1, topk_row, -1)
# Mask of the valid indices
mask = (indices_row != -1) & (indices_row < num_of_local_experts)

row_idx_offset = row_idx * num_of_local_experts
# Store to initialize
tl.store(multihot_indices_ptr + row_idx_offset + tl.arange(0, num_of_local_experts), 0)
tl.store(probs_in_multihot_ptr + row_idx_offset + tl.arange(0, num_of_local_experts), 0)
tl.store(position_map_ptr + row_idx_offset + tl.arange(0, num_of_local_experts), -1)
tl.store(multihot_indices_ptr + row_idx_offset + num_exp_row, 0, mask=num_exp_row_mask)
tl.store(probs_in_multihot_ptr + row_idx_offset + num_exp_row, 0, mask=num_exp_row_mask)
tl.store(position_map_ptr + row_idx_offset + num_exp_row, -1, mask=num_exp_row_mask)
# Use barrier to make sure the initialization is done
tl.debug_barrier()
# Store the indices and probs_in_indices
Expand All @@ -84,7 +97,9 @@ def _multihot_to_indices_kernel(
position_map_ptr,
probs_indices_ptr,
num_of_local_experts: tl.constexpr,
num_of_local_experts_next_power_of_2: tl.constexpr,
topk: tl.constexpr,
topk_next_power_of_2: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
'''
Expand Down Expand Up @@ -114,17 +129,27 @@ def _multihot_to_indices_kernel(
[0.1, 0.9]
]
'''
# Prepare the [0, topk) row
topk_row = tl.arange(0, topk_next_power_of_2)
topk_row = tl.where(topk_row < topk, topk_row, -1)
topk_row_mask = topk_row != -1
# Prepare the [0, num_of_local_experts) row
num_exp_row = tl.arange(0, num_of_local_experts_next_power_of_2)
num_exp_row = tl.where(num_exp_row < num_of_local_experts, num_exp_row, -1)
num_exp_row_mask = num_exp_row != -1

# Load a [1, num_of_local_experts] row from the local routing map
row_idx = tl.program_id(0)
ptr_offset = row_idx * num_of_local_experts + tl.arange(0, num_of_local_experts)
probs_in_multihot_row = tl.load(probs_in_multihot_ptr + ptr_offset)
ptr_offset = row_idx * num_of_local_experts + num_exp_row
probs_in_multihot_row = tl.load(probs_in_multihot_ptr + ptr_offset, mask=num_exp_row_mask)

# Get the original position of the valid value in the the indices
position_map_row = tl.load(position_map_ptr + ptr_offset)
position_map_row = tl.load(position_map_ptr + ptr_offset, mask=num_exp_row_mask)
position_map_row = tl.where(num_exp_row_mask, position_map_row, -1)
mask = position_map_row != -1

# Store to initialize
tl.store(probs_indices_ptr + row_idx * topk + tl.arange(0, topk), 0)
tl.store(probs_indices_ptr + row_idx * topk + topk_row, 0, mask=topk_row_mask)
# Use barrier to make sure the initialization is done
tl.debug_barrier()
# Restore the indices and probs_indices
Expand Down Expand Up @@ -169,6 +194,9 @@ def forward(ctx, indices, probs_indices, num_of_local_experts):
position_map = torch.empty(
(num_of_tokens, num_of_local_experts), dtype=torch.int32, device="cuda"
)
# Compute the next power of 2 for the topk and num_of_local_experts
topk_next_power_of_2 = 2 ** int(math.ceil(math.log2(topk)))
num_of_local_experts_next_power_of_2 = 2 ** int(math.ceil(math.log2(num_of_local_experts)))
grid = (num_of_tokens,)
_indices_to_multihot_kernel[grid](
indices,
Expand All @@ -177,7 +205,9 @@ def forward(ctx, indices, probs_indices, num_of_local_experts):
probs_in_multihot,
position_map,
num_of_local_experts,
num_of_local_experts_next_power_of_2,
topk,
topk_next_power_of_2,
BLOCK_SIZE=32, # use only 1 warp per block
num_warps=1,
)
Expand Down Expand Up @@ -211,6 +241,9 @@ def backward(ctx, grad_multihot_indices, grad_probs_in_multihot):
grad_probs_indices = torch.empty(
(num_of_tokens, topk), dtype=grad_probs_in_multihot.dtype, device="cuda"
)
# Compute the next power of 2 for the topk and num_of_local_experts
topk_next_power_of_2 = 2 ** int(math.ceil(math.log2(topk)))
num_of_local_experts_next_power_of_2 = 2 ** int(math.ceil(math.log2(num_of_local_experts)))

grid = (num_of_tokens,)
_multihot_to_indices_kernel[grid](
Expand All @@ -220,7 +253,9 @@ def backward(ctx, grad_multihot_indices, grad_probs_in_multihot):
position_map,
grad_probs_indices,
num_of_local_experts,
num_of_local_experts_next_power_of_2,
topk,
topk_next_power_of_2,
BLOCK_SIZE=32, # use only 1 warp per block
num_warps=1,
)
Expand Down