xx

Author: wangxiaoxin (A)

tests/sample/test_sampler.py

@@ -6,16 +6,16 @@
 
 from typing import Optional
 
-import pytest
 import torch
 
-from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p # noqa: F401
-from vllm.v1.sample.sampler import Sampler # noqa: F401
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p  # noqa: F401
+from vllm.v1.sample.sampler import Sampler  # noqa: F401
 
 # Set tolerance to 1 for quant ops
 DEFAULT_ATOL = 1e-3
 DEFAULT_RTOL = 1e-3
 
+
 def apply_min_p_new(
     logits: torch.Tensor,
     min_p: torch.Tensor,
@@ -28,14 +28,13 @@ def apply_min_p_new(
     # Convert logits to probability distribution
     probability_values = torch.nn.functional.softmax(logits, dim=-1)
     # Calculate maximum probabilities per sequence
-    max_probabilities = torch.amax(probability_values,
-                                    dim=-1,
-                                    keepdim=True)
+    max_probabilities = torch.amax(probability_values, dim=-1, keepdim=True)
     # Reshape min_p for broadcasting
     adjusted_min_p = min_p.unsqueeze(1) * max_probabilities
     # Identify valid tokens using threshold comparison
     # Apply mask using boolean indexing
-    logits = logits.masked_fill(probability_values < adjusted_min_p, -float('inf'))
+    logits = logits.masked_fill(probability_values < adjusted_min_p,
+                                -float('inf'))
     return logits
 
 
@@ -46,21 +45,23 @@ def apply_top_k_top_p_new(
 ) -> torch.Tensor:
     batch_size, vocab_size = logits.shape
     logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
-    
+
     # Apply top-k.
     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"))
-    
-    
+
     if p is not None:
         # Apply top-p.
         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)
+        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)
         logits_flatten = logits.flatten()
@@ -69,11 +70,12 @@ def apply_top_k_top_p_new(
         logits[valid_idx] = valid_logits
     return logits.reshape(batch_size, vocab_size)
 
+
 # test with leading dimension and merge seqlen and batch_size as num_tokens
 @torch.inference_mode()
 def test_apply_min_p(
 ) -> None:
-    logits = torch.randn((128,7168)).npu()
+    logits = torch.randn((128, 7168)).npu()
     min_p = torch.Tensor([0.01]).npu()
     logits_new = apply_min_p_new(logits, min_p)
     sampler = Sampler()
@@ -84,11 +86,12 @@ def test_apply_min_p(
                                atol=DEFAULT_ATOL,
                                rtol=DEFAULT_RTOL)
 
+
 # test with leading dimension and merge seqlen and batch_size as num_tokens
 @torch.inference_mode()
 def test_apply_top_k_top_p(
 ) -> None:
-    logits = torch.randn((128,7168)).npu()
+    logits = torch.randn((128, 7168)).npu()
     k = torch.Tensor([-1]).int().npu()
     p = torch.Tensor([1]).int().npu()
     logits_new = apply_top_k_top_p_new(logits, k, p)

vllm_ascend/patch/worker/patch_common/patch_sampler.py

@@ -61,7 +61,7 @@ def _apply_top_k_top_p(
     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,