refactor: support scenarios where top_p or top_k is None

linfeng-yuan · linfeng-yuan · commit d377ba3acf89 · 2025-04-28T13:36:29.000+08:00
Signed-off-by: linfeng-yuan &lt;1102311262@qq.com&gt;
diff --git a/vllm_ascend/sample/ops/ascend_topk_topp_sampler.py b/vllm_ascend/sample/ops/ascend_topk_topp_sampler.py
@@ -1,23 +1,12 @@
 from typing import Dict, Optional
 
 import torch
-import torch.nn as nn
-
 from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler, random_sample
-from vllm.logger import init_logger
-
-
-logger = init_logger(__name__)
 
 
 class AscendTopKTopPSampler(TopKTopPSampler):
 
-    def __init__(self):
-        super().__init__()
-        # TODO(linfeng): eliminate warning for FlashInfer here
-        self.forward = self.forward_npu
-
-    def forward_npu(
+    def forward_native(
         self,
         logits: torch.Tensor,
         generators: Dict[int, torch.Generator],
@@ -28,37 +17,48 @@ def forward_npu(
         logits = apply_top_k_top_p_npu(logits, k, p)
         probs = logits.softmax(dim=-1, dtype=torch.float32)
         return random_sample(probs, generators)
-    
+
 
 def apply_top_k_top_p_npu(
     logits: torch.Tensor,
     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)
diff --git a/vllm_ascend/sample/ops/penalties.py b/vllm_ascend/sample/ops/penalties.py
@@ -1,9 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
-
-from vllm.v1.sample.ops.penalties import _convert_to_tensors
 from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask
+from vllm.v1.sample.ops.penalties import _convert_to_tensors
 
 
 def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,
@@ -31,23 +30,25 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,
     output_bin_counts, output_mask = get_token_bin_counts_and_mask(
         output_tokens_tensor, vocab_size, num_seqs)
 
-    
     repetition_penalties = repetition_penalties.unsqueeze(dim=1).repeat(
         1, vocab_size)
-    
+
     # Avoid IndexPut operations in original apply_penalties function which are extremely time-consuming on NPU.
     sequence_mask = prompt_mask | output_mask
-    logits = torch.where(sequence_mask & torch.lt(logits, 0), logits * repetition_penalties,
-                            logits).to(logits.dtype)
-    logits = torch.where(sequence_mask & torch.ge(logits, 0), logits / repetition_penalties,
-                            logits).to(logits.dtype)
+    logits = torch.where(sequence_mask & torch.lt(logits, 0),
+                         logits * repetition_penalties,
+                         logits).to(logits.dtype)
+    logits = torch.where(sequence_mask & torch.ge(logits, 0),
+                         logits / repetition_penalties,
+                         logits).to(logits.dtype)
 
     # We follow the definition in OpenAI API.
     # Refer to https://platform.openai.com/docs/api-reference/parameter-details
     logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts
     logits -= presence_penalties.unsqueeze(dim=1) * output_mask
     return logits
 
+
 def apply_all_penalties(
     logits: torch.Tensor,
     prompt_token_ids: torch.Tensor,
@@ -64,4 +65,4 @@ def apply_all_penalties(
                                           logits.device)
     return apply_penalties(logits, prompt_token_ids, output_tokens_t,
                            presence_penalties, frequency_penalties,
-                           repetition_penalties)
+                           repetition_penalties)
diff --git a/vllm_ascend/sample/sampler.py b/vllm_ascend/sample/sampler.py
@@ -3,15 +3,13 @@
 from typing import Optional
 
 import torch
-from vllm.model_executor.layers.sampler import (Sampler, 
-                                                SamplerOutput, 
-                                                _apply_min_tokens_penalty, 
-                                                _apply_min_p, 
-                                                _sample, 
-                                                SampleResultArgsType, 
-                                                get_logprobs, 
-                                                _build_sampler_output)
+from vllm.model_executor.layers.sampler import (Sampler, SampleResultArgsType,
+                                                SamplerOutput, _apply_min_p,
+                                                _apply_min_tokens_penalty,
+                                                _build_sampler_output, _sample,
+                                                get_logprobs)
 from vllm.model_executor.sampling_metadata import SamplingMetadata
+
 from vllm_ascend.sample.ops.penalties import apply_penalties
 
 
@@ -61,7 +59,7 @@ def forward(
 
         if do_top_p_top_k:
             logits = _apply_top_k_top_p_npu(logits, sampling_tensors.top_ps,
-                                        sampling_tensors.top_ks)
+                                            sampling_tensors.top_ks)
 
         if do_min_p:
             logits = _apply_min_p(logits, sampling_tensors.min_ps)
@@ -83,21 +81,15 @@ def forward(
         )
 
         if self.include_gpu_probs_tensor:
-            # Since we will defer sampler result Pythonization,
-            # preserve GPU-side tensors in support of later
-            # deferred pythonization of logprobs
             assert maybe_sampled_tokens_tensor is not None
             on_device_tensors = (probs, logprobs, maybe_sampled_tokens_tensor)
         else:
-            # Since Pythonization has already happened, don't preserve
-            # GPU-side tensors.
             on_device_tensors = None
 
         # Get the logprobs query results.
         prompt_logprobs = None
         sample_logprobs = None
         if not sampling_metadata.skip_sampler_cpu_output:
-            # Pythonize logprobs now (GPU -> CPU); do not defer.
             assert not isinstance(maybe_deferred_sample_results,
                                   SampleResultArgsType)
             prompt_logprobs, sample_logprobs = get_logprobs(
@@ -121,10 +113,9 @@ def _apply_top_k_top_p_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
     batch_size, vocab_size = logits.shape
     logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
-    
+
     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"))
@@ -133,7 +124,10 @@ def _apply_top_k_top_p_npu(
     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()
diff --git a/vllm_ascend/sample/sampler_v1.py b/vllm_ascend/sample/sampler_v1.py
@@ -1,13 +1,11 @@
 import torch
-from vllm.v1.sample.sampler import Sampler
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.ops.penalties import apply_min_token_penalties
-from vllm.logger import init_logger
-from vllm_ascend.sample.ops.ascend_topk_topp_sampler import AscendTopKTopPSampler
-from vllm_ascend.sample.ops.penalties import apply_all_penalties
-
+from vllm.v1.sample.sampler import Sampler
 
-logger = init_logger(__name__)
+from vllm_ascend.sample.ops.ascend_topk_topp_sampler import \
+    AscendTopKTopPSampler
+from vllm_ascend.sample.ops.penalties import apply_all_penalties
 
 
 class AscendSampler(Sampler):
@@ -35,4 +33,4 @@ def apply_penalties(
                 sampling_metadata.repetition_penalties,
                 sampling_metadata.output_token_ids,
             )
-        return logits
+        return logits
diff --git a/vllm_ascend/worker/model_runner.py b/vllm_ascend/worker/model_runner.py
@@ -60,6 +60,7 @@
     _add_sampling_metadata_broadcastable_dict,
     _init_attn_metadata_from_tensor_dict,
     _init_sampling_metadata_from_tensor_dict)
+
 from vllm_ascend.sample.sampler import AscendSampler
 
 if TYPE_CHECKING:
@@ -821,12 +822,7 @@ def load_model(self) -> None:
         logger.info("Starting to load model %s...", self.model_config.model)
         with DeviceMemoryProfiler() as m:
             self.model = get_model(vllm_config=self.vllm_config)
-            # Same options with those in model_runner_v1.py
-            # option 1 
-            if hasattr(self.model, "sampler"):
-                self.model.sampler = AscendSampler()
-            # option 2
-            # self.model = NPUModelWrapperV1(model)
+            self.model.sampler = AscendSampler()
         self.model_memory_usage = m.consumed_memory
         logger.info("Loading model weights took %.4f GB",
                     self.model_memory_usage / float(2**30))
diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py
@@ -33,9 +33,7 @@
 from vllm.inputs import INPUT_REGISTRY
 from vllm.logger import logger
 from vllm.model_executor.layers.fused_moe import FusedMoE
-from vllm.model_executor.layers.sampler import sampler_output
 from vllm.model_executor.model_loader import get_model
-from vllm.model_executor.sampling_metadata import SamplingMetadata
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingType
@@ -813,11 +811,7 @@ def load_model(self) -> None:
 
         with DeviceMemoryProfiler() as m:  # noqa: SIM117
             self.model = get_model(vllm_config=self.vllm_config)
-            # option 1 
-            if hasattr(self.model, "sampler"):
-                self.model.sampler = AscendSampler()
-            # option 2
-            # self.model = NPUModelWrapperV1(model)
+            self.model.sampler = AscendSampler()
 
             if self.lora_config:
                 raise ValueError("LoRA model is not supported on NPU now.")
@@ -898,25 +892,3 @@ def get_kv_cache_spec(self) -> KVCacheSpec:
                     f"Unknown attention type: {attn_module.attn_type}")
 
         return kv_cache_spec
-
-# class NPUModelWrapperV1(nn.Module):
-
-#     def __init__(self, model: nn.Module):
-#         super().__init__()
-#         self._model = model
-#         self.sampler = AscendSampler()
-    
-#     def __getattr__(self, name):
-#         return getattr(self._model, name)
-
-#     def sample(
-#         self,
-#         logits: Optional[torch.Tensor],
-#         sampling_metadata: SamplingMetadata,
-#     ) -> Optional[SamplerOutput]:
-#         next_tokens = self.sampler(logits, sampling_metadata)
-#         return next_tokens
-    
-#     def forward():
-#         # necessary if using wrapper class
-#         pass
