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 5, 09:45 UTC
[Performance Improvement] Update `batched_count_greater_than` to handle batch size 1 without recompile (#38933)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Lucas Kabela committed
Apr 9, 2026 at 08:51 UTC
a8c6ee9b787d273916206a29b77feebadb80c368
3 files changed
+100
-1
tests/v1/sample/test_batched_count_greater_than.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Test that batched_count_greater_than does not trigger 0/1 specialization
4
+recompiles when batch_size varies."""
5
+
6
+import torch
7
+
8
+from vllm.platforms import current_platform
9
+from vllm.v1.sample.ops.logprobs import batched_count_greater_than
10
+from vllm.v1.sample.sampler import Sampler
11
+
12
+DEVICE = current_platform.device_type
13
+
14
+
15
+def test_batched_count_greater_than_correctness():
16
+ """Basic correctness: counts elements >= the corresponding value."""
17
+ x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device=DEVICE)
18
+ values = torch.tensor([[2.0], [5.0]], device=DEVICE)
19
+ result = batched_count_greater_than(x, values)
20
+ expected = torch.tensor([2, 2], device=DEVICE)
21
+ torch.testing.assert_close(result, expected)
22
+
23
+
24
+def test_gather_logprobs_no_recompile():
25
+ """Sampler.gather_logprobs with batch_size=1 then 2 must not recompile.
26
+
27
+ This guards against 0/1 specialization: dynamo normally specializes on
28
+ tensor sizes 0 and 1, causing a recompile when the size first exceeds 1.
29
+ The mark_unbacked calls in gather_logprobs prevent this.
30
+ """
31
+ torch._dynamo.reset()
32
+
33
+ compile_count = 0
34
+ orig_backend = current_platform.simple_compile_backend
35
+
36
+ def counting_backend(gm, example_inputs):
37
+ nonlocal compile_count
38
+ compile_count += 1
39
+ if orig_backend == "inductor":
40
+ return torch._inductor.compile(gm, example_inputs)
41
+ return gm
42
+
43
+ # Monkey-patch batched_count_greater_than with our counting backend
44
+ # so we can detect recompiles through the production code path.
45
+ import vllm.v1.sample.ops.logprobs as logprobs_module
46
+ import vllm.v1.sample.sampler as sampler_module
47
+
48
+ unwrapped = batched_count_greater_than._torchdynamo_orig_callable
49
+ patched = torch.compile(unwrapped, backend=counting_backend)
50
+ orig_fn = logprobs_module.batched_count_greater_than
51
+
52
+ logprobs_module.batched_count_greater_than = patched
53
+ sampler_module.batched_count_greater_than = patched
54
+
55
+ try:
56
+ vocab_size = 32
57
+ num_logprobs = 3
58
+
59
+ # Call 1: batch_size=1
60
+ logprobs1 = torch.randn(1, vocab_size, device=DEVICE)
61
+ token_ids1 = torch.randint(
62
+ 0, vocab_size, (1,), device=DEVICE, dtype=torch.int64
63
+ )
64
+ Sampler.gather_logprobs(logprobs1, num_logprobs, token_ids1)
65
+ assert compile_count == 1, f"Expected 1 compile, got {compile_count}"
66
+
67
+ # Call 2: batch_size=2 — should NOT recompile
68
+ logprobs2 = torch.randn(2, vocab_size, device=DEVICE)
69
+ token_ids2 = torch.randint(
70
+ 0, vocab_size, (2,), device=DEVICE, dtype=torch.int64
71
+ )
72
+ Sampler.gather_logprobs(logprobs2, num_logprobs, token_ids2)
73
+ assert compile_count == 1, (
74
+ f"Recompiled on batch_size 1->2 (0/1 specialization). "
75
+ f"Expected 1 compile, got {compile_count}"
76
+ )
77
+
78
+ # Call 3: batch_size=8 — should NOT recompile
79
+ logprobs3 = torch.randn(8, vocab_size, device=DEVICE)
80
+ token_ids3 = torch.randint(
81
+ 0, vocab_size, (8,), device=DEVICE, dtype=torch.int64
82
+ )
83
+ Sampler.gather_logprobs(logprobs3, num_logprobs, token_ids3)
84
+ assert compile_count == 1, (
85
+ f"Recompiled on batch_size change. Expected 1 compile, got {compile_count}"
86
+ )
87
+ finally:
88
+ # Restore original function
89
+ logprobs_module.batched_count_greater_than = orig_fn
90
+ sampler_module.batched_count_greater_than = orig_fn
91
+ torch._dynamo.reset()
vllm/v1/sample/ops/logprobs.py
+3
-1
@@ -7,7 +7,7 @@ import torch
7
from vllm.platforms import current_platform
8
9
10
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
10
+@torch.compile(backend=current_platform.simple_compile_backend)
11
def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.Tensor:
12
"""
13
Counts elements in each row of x that are greater than the corresponding
@@ -22,4 +22,6 @@ def batched_count_greater_than(x: torch.Tensor, values: torch.Tensor) -> torch.T
22
Returns:
23
torch.Tensor: A 1D tensor of shape (batch_size,) with the counts.
24
"""
25
+ torch._check(x.shape[0] >= 1)
26
+ torch._check(x.shape[0] == values.shape[0])
27
return (x >= values).sum(-1)
vllm/v1/sample/sampler.py
+6
@@ -324,6 +324,12 @@ class Sampler(nn.Module):
324
token_logprobs = logprobs.gather(-1, token_ids)
325
326
# Compute the ranks of the actual token.
327
+ # Avoid 0/1 specialization recompile on the batch dimension
328
+ # of the compiled batched_count_greater_than. mark_unbacked makes
329
+ # the size fully symbolic so dynamo doesn't specialize when
330
+ # batch_size transitions from 1 to >=2.
331
+ torch._dynamo.decorators.mark_unbacked(logprobs, 0)
332
+ torch._dynamo.decorators.mark_unbacked(token_logprobs, 0)
333
token_ranks = batched_count_greater_than(logprobs, token_logprobs)
334
335
# Concatenate together with the topk.