fix lora tests

Author: gnovack

tests/lora/test_fused_moe_lora_kernel.py

@@ -135,6 +135,11 @@ def use_fused_moe_lora_kernel(
     expert_ids = torch.empty((max_loras * max_num_m_blocks,), dtype=torch.int32)
     num_tokens_post_padded = torch.empty((max_loras,), dtype=torch.int32)
 
+    num_tokens_per_lora = torch.ones(max_loras+1, dtype=torch.int32)
+    adapter_enabled = torch.ones(max_loras+1, dtype=torch.int32)
+    lora_ids = torch.arange(1,max_loras+1, dtype=torch.int32)
+
+
     # call kernel
     ops.moe_lora_align_block_size(
         topk_ids,
@@ -147,6 +152,8 @@ def use_fused_moe_lora_kernel(
         sorted_token_ids,
         expert_ids,
         num_tokens_post_padded,
+        num_tokens_per_lora,
+        adapter_enabled,
     )
 
     config = {
@@ -171,6 +178,8 @@ def use_fused_moe_lora_kernel(
         num_tokens_post_padded,
         max_lora_rank,
         top_k_num,
+        lora_ids,
+        adapter_enabled,
         config["BLOCK_SIZE_M"],
         config["BLOCK_SIZE_N"],
         config["BLOCK_SIZE_K"],

tests/lora/test_moe_lora_align_sum.py

@@ -60,15 +60,13 @@ def test_moe_lora_align_block_size(
         (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
     )
     num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
-    num_tokens_per_lora = torch.ones((max_loras,), dtype=torch.int32, device="cuda")
-    adapter_enabled = torch.ones((max_loras,), dtype=torch.int32, device="cuda")
+    num_tokens_per_lora = torch.ones((max_loras+1,), dtype=torch.int32, device="cuda")
+    adapter_enabled = torch.ones((max_loras+1,), dtype=torch.int32, device="cuda")
 
     # call kernel
     ops.moe_lora_align_block_size(
         topk_ids,
         token_lora_mapping,
-        num_tokens_per_lora,
-        adapter_enabled,
         num_experts,
         block_size,
         max_loras,
@@ -77,6 +75,8 @@ def test_moe_lora_align_block_size(
         sorted_token_ids,
         expert_ids,
         num_tokens_post_pad,
+        num_tokens_per_lora,
+        adapter_enabled,
     )
 
     # verify values
@@ -91,73 +91,3 @@ def test_moe_lora_align_block_size(
                 expert_id = expert_ids[lora_idx][token_idx]
                 assert torch.all(topk_ids.view(-1)[indices] == expert_id)
 
-@pytest.mark.parametrize("num_tokens", [4096])
-@pytest.mark.parametrize("topk_num", [6])
-@pytest.mark.parametrize("num_experts", [64])
-@pytest.mark.parametrize("max_loras", [2])
-@pytest.mark.parametrize("block_size", [16])
-@pytest.mark.parametrize("adapter_enabled", [[0,1],[0,0]])
-def test_moe_lora_align_block_size_early_exit(
-    num_tokens, topk_num, num_experts, max_loras, block_size, adapter_enabled
-):
-
-    # sample data
-    random.seed(1)
-    topk_ids, token_lora_mapping = sample_data(
-        num_experts, max_loras, num_tokens, topk_num
-    )
-
-    # compute paddings
-    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
-    max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
-    max_num_m_blocks = CEILDIV(max_num_tokens_padded, block_size)
-
-    # init output tensors
-    sorted_token_ids = torch.full(
-        (max_loras * max_num_tokens_padded,),
-        topk_ids.numel(),
-        dtype=torch.int32,
-        device="cuda",
-    )
-    expert_ids = torch.full(
-        (max_loras * max_num_m_blocks,), num_experts, dtype=torch.int32, device="cuda"
-    )
-    num_tokens_post_pad = torch.zeros((max_loras,), dtype=torch.int32, device="cuda")
-    
-    num_tokens_per_lora = torch.ones((max_loras,), dtype=torch.int32, device="cuda")
-    adapter_enabled = torch.tensor(adapter_enabled, dtype=torch.int32, device="cuda")
-
-    # call kernel
-    ops.moe_lora_align_block_size(
-        topk_ids,
-        token_lora_mapping,
-        num_tokens_per_lora,
-        adapter_enabled,
-        num_experts,
-        block_size,
-        max_loras,
-        sorted_token_ids,
-        expert_ids,
-        num_tokens_post_pad,
-    )
-
-    # verify values
-    expert_ids = expert_ids.view(max_loras, -1)
-    sorted_token_ids = sorted_token_ids.view(max_loras, -1, block_size)
-
-    for lora_idx in range(max_loras):
-
-        # assert not operation was performed
-        if adapter_enabled[lora_idx].item() == 0:
-            assert torch.all(sorted_token_ids[lora_idx] == topk_ids.numel())
-        else:
-            for token_idx in range(sorted_token_ids.size(1)):
-                block = sorted_token_ids[lora_idx][token_idx]
-                indices = block[block != topk_ids.numel()]
-                if indices.numel() > 0:
-                    expert_id = expert_ids[lora_idx][token_idx]
-                    assert torch.all(topk_ids.view(-1)[indices] == expert_id)
-
-
-if __name__ == "__main__":
-    pytest.main([__file__])

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

@@ -337,12 +337,12 @@ def _fused_moe_lora_fake(
     num_tokens_post_padded: torch.Tensor,
     max_lora_rank: int,
     top_k_num: int,
+    lora_ids: torch.Tensor,
+    adapter_enabled: torch.Tensor,
     block_size_m: int,
     block_size_n: int,
     block_size_k: int,
     group_size_m: int,
-    lora_ids: torch.Tensor,
-    adapter_enabled: torch.Tensor,
     mul_routed_weight: bool = False,
 ) -> None:
     return