Fix moe align kernel test (#8531)

ispobock · web-flow · commit 5973675bc30d · 2025-07-29T11:03:02.000-07:00
diff --git a/sgl-kernel/tests/test_moe_align.py b/sgl-kernel/tests/test_moe_align.py
@@ -229,7 +229,7 @@ def test_moe_align_block_size_compare_implementations(
     matching_indices = torch.where(expert_ids_cuda == expert_idx)[0]
     block_sorted_start = matching_indices[0].item() * block_size
     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[

Original file line number	Diff line number	Diff line change
`@@ -229,7 +229,7 @@ def test_moe_align_block_size_compare_implementations(`
`229`	`229`	`matching_indices = torch.where(expert_ids_cuda == expert_idx)[0]`
`230`	`230`	`block_sorted_start = matching_indices[0].item() * block_size`
`231`	`231`	`block_sorted_end = min(`
`232`		`- (matching_indices[-1].item() + 1) * block_size, max_num_tokens_padded`
	`232`	`+ (matching_indices[-1].item() + 1) * block_size, num_tokens_post_pad_cuda.item()`
`233`	`233`	`)`
`234`	`234`
`235`	`235`	`selected_sorted_ids_cuda = sorted_ids_cuda[`