add ut

Author: RunningLeon

tests/pytorch/paging/test_block_trie.py

@@ -2,7 +2,7 @@
 import pytest
 
 from lmdeploy.pytorch.config import CacheConfig
-from lmdeploy.pytorch.messages import SchedulerSession, SequenceManager, SequenceMeta
+from lmdeploy.pytorch.messages import SamplingParam, SchedulerSession, SequenceManager, SequenceMeta
 from lmdeploy.pytorch.paging.block_manager import build_block_manager
 from lmdeploy.pytorch.paging.block_trie import BlockTrie
 
@@ -37,13 +37,55 @@ def block_mgr(self, cache_config):
     def block_trie(self, cache_config, block_mgr):
         yield BlockTrie(cache_config, block_mgr)
 
+    @pytest.fixture
+    def num_moe_layers(self):
+        yield 4
+
+    @pytest.fixture
+    def experts_topk(self):
+        yield 4
+
     @pytest.fixture
     def seq_manager(self, block_size):
         from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
         strategy = ARSequenceStrategy()
         seq_meta = SequenceMeta(block_size, strategy=strategy)
         yield SequenceManager(seq_meta)
 
+    def test_with_routed_experts(self, block_trie, block_mgr, seq_manager, num_moe_layers, experts_topk):
+
+        def _get_routed_experts(size, value):
+            return np.full((size, num_moe_layers, experts_topk), value, dtype=np.int32)
+
+        sess = SchedulerSession(0, seq_manager)
+        block_size = sess.seq_meta.block_size
+        token_ids = ([1] * block_size + [2] * block_size)
+        all_routed_experts = [_get_routed_experts(block_size, 1), _get_routed_experts(block_size, 2)]
+        token_ids += [3] * (block_size // 2)
+        all_routed_experts += [_get_routed_experts(block_size // 2, 3)]
+        seq = sess.add_sequence(token_ids, sampling_param=SamplingParam(return_routed_experts=True))
+        all_routed_experts += [_get_routed_experts(block_size - 1, 4)]
+        routed_experts = np.concatenate(all_routed_experts, axis=0)
+        seq.update_token_ids([4] * block_size, routed_experts=routed_experts)
+
+        # test allocate
+        block_mgr.allocate(seq)
+        block_trie.allocate(seq)
+        node = getattr(seq.logical_blocks, 'last_shared_node', None)
+        assert node is not None
+        assert node.routed_experts is not None
+        target_routed_experts = np.concatenate(
+            [_get_routed_experts(block_size // 2, 3),
+             _get_routed_experts(block_size // 2, 4)], axis=0)
+        assert np.array_equal(node.routed_experts, target_routed_experts)
+
+        # test match
+        seq_query = sess.add_sequence(token_ids, sampling_param=SamplingParam(return_routed_experts=True))
+        block_trie.match(seq_query)
+        assert seq_query.all_routed_experts is not None
+        assert len(seq_query.all_routed_experts) == block_size * 2
+        assert np.array_equal(seq_query.all_routed_experts.get_real(), np.concatenate(all_routed_experts[:2], axis=0))
+
     def test_allocate(self, block_trie, block_mgr, seq_manager):
         allocator = block_trie.allocator
         sess = SchedulerSession(0, seq_manager)