Unverified Commit 5973675b authored by Ke Bao's avatar Ke Bao Committed by GitHub
Browse files

Fix moe align kernel test (#8531)

parent 4d16c88b
...@@ -229,7 +229,7 @@ def test_moe_align_block_size_compare_implementations( ...@@ -229,7 +229,7 @@ def test_moe_align_block_size_compare_implementations(
matching_indices = torch.where(expert_ids_cuda == expert_idx)[0] matching_indices = torch.where(expert_ids_cuda == expert_idx)[0]
block_sorted_start = matching_indices[0].item() * block_size block_sorted_start = matching_indices[0].item() * block_size
block_sorted_end = min( block_sorted_end = min(
(matching_indices[-1].item() + 1) * block_size, max_num_tokens_padded (matching_indices[-1].item() + 1) * block_size, num_tokens_post_pad_cuda.item()
) )
selected_sorted_ids_cuda = sorted_ids_cuda[ selected_sorted_ids_cuda = sorted_ids_cuda[
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment