fix lora masking

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

Author: gnovack

csrc/moe/moe_align_sum_kernels.cu

@@ -246,10 +246,10 @@ __device__ void _moe_align_block_size_small_batch_expert(
     int32_t rank_post_pad =
         tokens_cnts[threadIdx.x * num_experts + expert_id] + cumsum[expert_id];
 
-    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
-    sorted_token_ids[sorted_token_ids_offset + rank_post_pad] +=
-        ((i - numel) * mask);
-    tokens_cnts[threadIdx.x * num_experts + expert_id] += mask;
+    if (token_mask == nullptr || token_mask[i / topk_num]) {
+      sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i;
+      ++tokens_cnts[threadIdx.x * num_experts + expert_id];
+    }
   }
 }
 
@@ -268,11 +268,12 @@ __device__ void _count_and_sort_expert_tokens(
       continue;
     }
 
-    int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num];
-    int32_t rank_post_pad = atomicAdd(
-        &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], mask);
-    sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] +=
-        ((i - numel) * mask);
+    if (token_mask == nullptr || token_mask[i / topk_num]) {
+      int32_t rank_post_pad = atomicAdd(
+          &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1);
+      sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] =
+          i;
+    }
   }
 }