A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[Minor][Spec Decode] Remove compiled_softmax (#15416)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Mar 24, 2025 at 19:09 UTC
911c8eb0000b1f9d1fef99ac9e209f83d801bd0a
2 files changed
+1
-33
vllm/v1/sample/ops/utils.py
deleted
-30
@@ -1,30 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-from typing import Union
3
-
4
-import torch
5
-
6
-
7
-def compiled_softmax(
8
- logits: torch.Tensor,
9
- temperature: Union[float, torch.Tensor] = 1.0,
10
-) -> torch.Tensor:
11
- """Faster softmax kernel generated by torch.compile.
12
-
13
- Args:
14
- logits: [n, vocab_size]
15
- temperature: [n] or float
16
- """
17
- # NOTE(woosuk): Avoid recompilation by marking the first dim as dynamic.
18
- torch._dynamo.mark_dynamic(logits, index=0)
19
- if isinstance(temperature, torch.Tensor):
20
- torch._dynamo.mark_dynamic(temperature, index=0)
21
- return _softmax(logits, temperature)
22
-
23
-
24
-@torch.compile
25
-def _softmax(
26
- logits: torch.Tensor,
27
- temperature: Union[float, torch.Tensor],
28
-) -> torch.Tensor:
29
- logits = logits / temperature
30
- return torch.softmax(logits, dim=-1, dtype=torch.float32)
vllm/v1/sample/rejection_sampler.py
+1
-3
@@ -9,7 +9,6 @@ import triton.language as tl
9
from vllm.logger import init_logger
10
from vllm.v1.sample.metadata import SamplingMetadata
11
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
12
-from vllm.v1.sample.ops.utils import compiled_softmax
12
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
13
14
logger = init_logger(__name__)
@@ -275,8 +274,7 @@ def compute_probs(
274
# NOTE(woosuk): `apply_top_k_top_p` uses sorting to calculate the mask,
275
# which is slow for large vocab sizes. This may cause performance issues.
276
logits = apply_top_k_top_p(logits, top_k, top_p)
278
-
279
- output_prob = compiled_softmax(logits)
277
+ output_prob = logits.softmax(dim=-1, dtype=torch.float32)
278
return output_prob
279
280