fix tests

Signed-off-by: gnovack <[email protected]>

Author: gnovack

csrc/moe/moe_lora_align_sum_kernels.cu

@@ -28,8 +28,8 @@ __global__ void moe_lora_align_sum_kernel(
     int64_t block_size, int num_experts, int max_loras, size_t numel,
     int max_num_tokens_padded, int max_num_m_blocks,
     int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids,
-    int topk_num, int32_t* total_tokens_post_pad, int32_t* num_tokens_per_lora,
-    int32_t* adapter_enabled, int32_t* lora_ids) {
+    int topk_num, int32_t* total_tokens_post_pad, int32_t* adapter_enabled,
+    int32_t* lora_ids) {
   const size_t tokens_per_thread = div_ceil(numel, blockDim.x);
   const size_t start_idx = threadIdx.x * tokens_per_thread;
 
@@ -131,8 +131,8 @@ void moe_lora_align_block_size(
     int64_t num_experts, int64_t block_size, int64_t max_loras,
     int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
     torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
-    torch::Tensor num_tokens_post_pad, torch::Tensor num_tokens_per_lora,
-    torch::Tensor adapter_enabled, torch::Tensor lora_ids) {
+    torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
+    torch::Tensor lora_ids) {
   const int topk_num = topk_ids.size(1);
 
   TORCH_CHECK(block_size > 0, "block_size should be greater than 0. ");
@@ -169,7 +169,6 @@ void moe_lora_align_block_size(
             max_num_m_blocks, sorted_token_ids.data_ptr<int32_t>(),
             expert_ids.data_ptr<int32_t>(), topk_num,
             num_tokens_post_pad.data_ptr<int32_t>(),
-            num_tokens_per_lora.data_ptr<int32_t>(),
             adapter_enabled.data_ptr<int32_t>(), lora_ids.data_ptr<int32_t>());
       });
 }

csrc/moe/moe_ops.h

@@ -25,8 +25,8 @@ void moe_lora_align_block_size(
     int64_t num_experts, int64_t block_size, int64_t max_loras,
     int64_t max_num_tokens_padded, int64_t max_num_m_blocks,
     torch::Tensor sorted_token_ids, torch::Tensor expert_ids,
-    torch::Tensor num_tokens_post_pad, torch::Tensor num_tokens_per_lora,
-    torch::Tensor adapter_enabled, torch::Tensor lora_ids);
+    torch::Tensor num_tokens_post_pad, torch::Tensor adapter_enabled,
+    torch::Tensor lora_ids);
 #ifndef USE_ROCM
 torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
                              torch::Tensor b_qweight, torch::Tensor b_scales,

csrc/moe/torch_bindings.cpp

@@ -45,7 +45,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
       "                     Tensor !sorted_token_ids,"
       "                     Tensor !experts_ids,"
       "                     Tensor !num_tokens_post_pad,"
-      "                     Tensor !num_tokens_per_lora,"
       "                     Tensor !adapter_enabled,"
       "                     Tensor !lora_ids) -> () ");
   m.impl("moe_lora_align_block_size", torch::kCUDA, &moe_lora_align_block_size);

tests/lora/test_fused_moe_lora_kernel.py

@@ -134,9 +134,8 @@ 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)
+    lora_ids = torch.arange(max_loras + 2, dtype=torch.int32)
 
     # call kernel
     ops.moe_lora_align_block_size(
@@ -150,8 +149,8 @@ def use_fused_moe_lora_kernel(
         sorted_token_ids,
         expert_ids,
         num_tokens_post_padded,
-        num_tokens_per_lora,
         adapter_enabled,
+        lora_ids,
     )
 
     config = {

tests/lora/test_moe_lora_align_sum.py

@@ -60,8 +60,8 @@ 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 + 1,), dtype=torch.int32, device="cuda")
     adapter_enabled = torch.ones((max_loras + 1,), dtype=torch.int32, device="cuda")
+    lora_ids = torch.arange(max_loras + 2, dtype=torch.int32, device="cuda")
 
     # call kernel
     ops.moe_lora_align_block_size(
@@ -75,8 +75,8 @@ def test_moe_lora_align_block_size(
         sorted_token_ids,
         expert_ids,
         num_tokens_post_pad,
-        num_tokens_per_lora,
         adapter_enabled,
+        lora_ids,
     )
 
     # verify values

tests/lora/test_olmoe_tp.py

@@ -60,7 +60,6 @@ def generate_and_test(
             for i in lora_id
         ]
 
-    print(f"Sending lora req {lora_request}")
     sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
     outputs = llm.generate(prompts, sampling_params, lora_request=lora_request)
     # Print the outputs.
@@ -78,7 +77,6 @@ def generate_and_test(
             if req_lora_id is not None
             else EXPECTED_BASE_MODEL_OUTPUT[i]
         )
-
         assert generated_texts[i].startswith(expected_output)
 
 
@@ -100,8 +98,6 @@ def test_olmoe_lora(olmoe_lora_files):
 
 
 def test_olmoe_lora_base_model(olmoe_lora_files):
-    # We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
-    # Otherwise, the lora-test will fail due to CUDA OOM.
     llm = vllm.LLM(
         MODEL_PATH,
         max_model_len=1024,

vllm/_custom_ops.py

@@ -1806,7 +1806,6 @@ def moe_lora_align_block_size(
     sorted_token_ids: torch.Tensor,
     experts_ids: torch.Tensor,
     num_tokens_post_pad: torch.Tensor,
-    num_tokens_per_lora: torch.Tensor,
     adapter_enabled: torch.Tensor,
     lora_ids: torch.Tensor,
 ) -> None:
@@ -1821,7 +1820,6 @@ def moe_lora_align_block_size(
         sorted_token_ids,
         experts_ids,
         num_tokens_post_pad,
-        num_tokens_per_lora,
         adapter_enabled,
         lora_ids,
     )

vllm/lora/layers/fused_moe.py

@@ -108,11 +108,6 @@ def wrapper(*args, **kwargs):
                     block_shape=layer.quant_method.moe_quant_config.block_shape,
                 )
 
-                (_, _, num_tokens_per_lora, _, _, _) = (
-                    self.punica_wrapper.token_mapping_meta.meta_args(
-                        hidden_states.size(0)
-                    )
-                )
                 max_loras = self.w1_lora_a_stacked.shape[0]
                 config = get_config_func(M)
                 (
@@ -125,7 +120,6 @@ def wrapper(*args, **kwargs):
                     config["BLOCK_SIZE_M"],
                     global_num_experts,
                     max_loras,
-                    num_tokens_per_lora,
                     self.adapter_enabled,
                     expert_map,
                 )

vllm/lora/punica_wrapper/punica_gpu.py

@@ -305,7 +305,6 @@ def moe_lora_align_block_size(
         block_size: int,
         num_experts: int,
         max_loras: int,
-        num_tokens_per_lora: torch.Tensor,
         adapter_enabled: torch.Tensor,
         expert_map: torch.Tensor | None = None,
         pad_sorted_ids: bool = False,
@@ -348,7 +347,6 @@ def moe_lora_align_block_size(
             sorted_ids,
             expert_ids,
             num_tokens_post_pad,
-            num_tokens_per_lora,
             adapter_enabled,
             lora_ids,
         )