68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
import unittest
|
|
|
|
import torch
|
|
|
|
from sglang.kernels.ops.moe.ep_moe_kernels import zero_experts_compute_triton
|
|
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
|
|
register_amd_ci(est_time=5, stage="stage-b", runner_config="1-gpu-small-amd")
|
|
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
|
|
class TestZeroExpertsComputeTriton(CustomTestCase):
|
|
def test_zero_expert_routes_keep_valid_ids(self):
|
|
num_experts = 4
|
|
hidden_states = torch.arange(
|
|
2 * 512, dtype=torch.float32, device="cuda"
|
|
).reshape(2, 512)
|
|
expert_indices = torch.tensor(
|
|
[[0, 4, 1], [5, 2, 6]], dtype=torch.int64, device="cuda"
|
|
)
|
|
expert_scales = torch.tensor(
|
|
[[0.1, 0.3, 0.2], [0.4, 0.5, 0.6]],
|
|
dtype=torch.float32,
|
|
device="cuda",
|
|
)
|
|
original_indices = expert_indices.clone()
|
|
original_scales = expert_scales.clone()
|
|
|
|
output = zero_experts_compute_triton(
|
|
expert_indices=expert_indices,
|
|
expert_scales=expert_scales,
|
|
num_experts=num_experts,
|
|
zero_expert_type="identity",
|
|
hidden_states=hidden_states,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
zero_expert_mask = original_indices >= num_experts
|
|
normal_expert_mask = ~zero_expert_mask
|
|
self.assertTrue(torch.all(expert_indices >= 0).item())
|
|
torch.testing.assert_close(
|
|
expert_indices[normal_expert_mask],
|
|
original_indices[normal_expert_mask],
|
|
)
|
|
torch.testing.assert_close(
|
|
expert_indices[zero_expert_mask],
|
|
torch.zeros_like(expert_indices[zero_expert_mask]),
|
|
)
|
|
torch.testing.assert_close(
|
|
expert_scales[normal_expert_mask],
|
|
original_scales[normal_expert_mask],
|
|
)
|
|
torch.testing.assert_close(
|
|
expert_scales[zero_expert_mask],
|
|
torch.zeros_like(expert_scales[zero_expert_mask]),
|
|
)
|
|
|
|
expected_scale = (
|
|
original_scales * zero_expert_mask.to(original_scales.dtype)
|
|
).sum(dim=-1, keepdim=True)
|
|
torch.testing.assert_close(output, hidden_states * expected_scale)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|