refactor: support scenarios where top_p or top_k is None

Signed-off-by: linfeng-yuan <1102311262@qq.com>

Author: linfeng-yuan

vllm_ascend/sample/ops/ascend_topk_topp_sampler.py

@@ -35,30 +35,36 @@ def apply_top_k_top_p_npu(
     k: Optional[torch.Tensor],
     p: Optional[torch.Tensor],
 ) -> torch.Tensor:
-    """Apply top-k and top-p optimized for NPU.
-
-    This algorithm avoids using torch.scatter which is time-consuming on NPU.
-    """
-    # TODO(linfeng): consider the case taht either p or k is applied
+    """Apply top-k and/or top-p optimized for NPU."""
     if k is None and p is None:
         return logits
+
     batch_size, vocab_size = logits.shape
+    device = logits.device
     logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
+    if k is not None:
+        safe_k = torch.clamp(k, min=1, max=vocab_size)
+        boundary_idx = (vocab_size - safe_k).unsqueeze(1)  
+        boundary = logits_sort.gather(1, boundary_idx)             
+        top_k_mask = logits_sort < boundary                       
+        logits_sort = logits_sort.masked_fill(top_k_mask, -float("inf"))
+    else:
+        top_k_mask = torch.zeros_like(logits_sort, dtype=torch.bool)
 
-    boundary = logits_sort.gather(1, (vocab_size - k).unsqueeze(dim=1))
-    top_k_mask = logits_sort < boundary
-    logits_sort.masked_fill_(top_k_mask, -float("inf"))
-    cutoff = top_k_mask.sum(dim=-1).min()
-    probs_sort = logits_sort.softmax(dim=-1)[:, cutoff:]
-    probs_sum = probs_sort.cumsum(dim=-1)
-    top_p_mask = probs_sum > 1 - p.unsqueeze(dim=1)
-    top_p_mask[:, -1] = True
-    strides = torch.arange(0, batch_size*vocab_size, vocab_size, device=logits.device)
-    flatten_idx = logits_idx[:, cutoff:] + strides.unsqueeze(dim=1)
-    valid_idx = torch.masked_select(flatten_idx, top_p_mask)
+    cutoffs = top_k_mask.sum(dim=-1)
+    strides = torch.arange(0, batch_size*vocab_size, vocab_size, device=device).unsqueeze(1)
+    if p is not None:
+        global_cutoff = cutoffs.min()
+        active_part = logits_idx[:, global_cutoff:]
+        probs_sort = logits_sort[:, global_cutoff:].softmax(dim=-1)
+        cumprob = probs_sort.cumsum(dim=-1)
+        top_p_mask = (cumprob <= (1 - p.unsqueeze(1))) | (torch.arange(probs_sort.size(1), device=device) == probs_sort.size(1)-1)
+    else:
+        active_part = logits_idx
+        top_p_mask = torch.arange(vocab_size, device=device).expand(batch_size, -1) >= cutoffs.unsqueeze(1)
 
+    valid_idx = (active_part + strides).masked_select(top_p_mask)
     logits_flatten = logits.flatten()
-    valid_logits = torch.index_select(logits_flatten, 0, valid_idx)
-    logits = torch.empty_like(logits_flatten).fill_(-float("inf"))
-    logits[valid_idx] = valid_logits
-    return logits.reshape(batch_size, vocab_size)
+    output = torch.full_like(logits_flatten, -float('inf'))
+    output[valid_idx] = logits_flatten[valid_idx]
+    return output.reshape(batch_size, vocab_size)