File size: 2,337 Bytes
4a28d4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import pytest
import torch
class TestSiluAndMul:
@pytest.fixture
def seqlen(self, request):
yield request.param
@pytest.fixture
def feat_size(self, request):
yield request.param
@pytest.fixture
def x(self, seqlen, feat_size):
yield torch.rand(seqlen, feat_size, dtype=torch.float16, device='cuda')
@pytest.fixture
def gt(self, x):
gate, up = x.chunk(2, -1)
gate = torch.nn.functional.silu(gate)
yield gate * up
@pytest.mark.parametrize('seqlen', [65536, 256], indirect=True)
@pytest.mark.parametrize('feat_size', [4096, 768], indirect=True)
def test_silu_and_mul(self, x, gt):
from lmdeploy.pytorch.kernels.cuda.activation import silu_and_mul
out = silu_and_mul(x)
torch.testing.assert_close(out, gt)
class TestSiluAndMulMoEEP:
@pytest.fixture
def num_experts(self, request):
yield request.param
@pytest.fixture
def seqlen(self, request):
yield request.param
@pytest.fixture
def feat_size(self, request):
yield request.param
@pytest.fixture
def dtype(self):
yield torch.float16
@pytest.fixture
def x(self, num_experts, seqlen, feat_size, dtype):
yield torch.rand(num_experts, seqlen, feat_size, dtype=dtype, device='cuda')
@pytest.fixture
def mask_m(self, num_experts, seqlen):
mask_m = torch.randint(0, seqlen, (num_experts, ), device='cuda')
yield mask_m
@pytest.fixture
def elem_mask(self, mask_m, seqlen):
elem_mask = torch.arange(seqlen, device='cuda').unsqueeze(0) < mask_m.unsqueeze(1)
yield elem_mask[..., None]
@pytest.fixture
def gt(self, x):
gate, up = x.chunk(2, -1)
gate = torch.nn.functional.silu(gate)
yield gate * up
@pytest.mark.parametrize('num_experts', [4], indirect=True)
@pytest.mark.parametrize('seqlen', [1024], indirect=True)
@pytest.mark.parametrize('feat_size', [4096, 768], indirect=True)
def test_silu_and_mul(self, x, mask_m, elem_mask, gt):
from lmdeploy.pytorch.kernels.cuda.activation import silu_and_mul_moe_ep
out = silu_and_mul_moe_ep(x, mask_m)
out.masked_fill_(~elem_mask, 0.0)
gt.masked_fill_(~elem_mask, 0.0)
torch.testing.assert_close(out, gt)
|