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 13, 15:45 UTC
[CPU]Avoid repeated random sample compile (#28260)
Signed-off-by: Zhang Xiangze <Xiangze.Zhang@arm.com>
Zhang Xiangze committed
Nov 7, 2025 at 19:03 UTC
7bdb42b2f22f14bf450e084b8f9938f598c08f9b
1 file changed
+10
-9
vllm/v1/sample/ops/topk_topp_sampler.py
+10
-9
@@ -127,15 +127,6 @@ class TopKTopPSampler(nn.Module):
127
elif self.logprobs_mode == "processed_logprobs":
128
logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
129
130
- # Note: this is a workaround for
131
- # https://github.com/pytorch/pytorch/pull/151218
132
- @torch.compile(dynamic=True)
133
- def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor:
134
- probs = logits.softmax(dim=-1, dtype=torch.float32)
135
- q = torch.empty_like(probs)
136
- q.exponential_()
137
- return probs.div(q).argmax(dim=-1).view(-1)
138
-
130
if len(generators) != logits.shape[0]:
131
return compiled_random_sample(logits), logits_to_return
132
else:
@@ -148,6 +139,16 @@ class TopKTopPSampler(nn.Module):
139
return probs.div_(q).argmax(dim=-1).view(-1), logits_to_return
140
141
142
+# Note: this is a workaround for
143
+# https://github.com/pytorch/pytorch/pull/151218
144
+@torch.compile(dynamic=True)
145
+def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor:
146
+ probs = logits.softmax(dim=-1, dtype=torch.float32)
147
+ q = torch.empty_like(probs)
148
+ q.exponential_()
149
+ return probs.div(q).argmax(dim=-1).view(-1)
150
+
151
+
152
def apply_top_k_top_p(
153
logits: torch.Tensor,
154
k: torch.Tensor | None,