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, 03:30 UTC
[BUG] Fix #20484. Support empty sequence in cuda penalty kernel (#20491)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@centml.ai>
Vadim Gimpelson committed
Jul 6, 2025 at 06:38 UTC
f73d02aadcc0c7f7b300528b32a92dca08e82138
2 files changed
+50
csrc/sampler.cu
+2
@@ -59,6 +59,8 @@ void apply_repetition_penalties_(
59
int vocab_size = logits.size(-1);
60
int num_seqs = logits.size(0);
61
62
+ if (num_seqs == 0) return;
63
+
64
// Get number of SMs on the current device
65
int sms = 0;
66
cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount,
tests/kernels/test_apply_repetition_penalties.py
+48
@@ -75,3 +75,51 @@ def test_apply_repetition_penalties(
75
# Test the operator by applying the opcheck utility
76
opcheck(torch.ops._C.apply_repetition_penalties_,
77
(logits.clone(), prompt_mask, output_mask, repetition_penalties))
78
+
79
+
80
+@pytest.mark.skipif(not current_platform.is_cuda(),
81
+ reason="This test for checking CUDA kernel")
82
+@torch.inference_mode()
83
+def test_apply_repetition_penalties_zero_seqs() -> None:
84
+ """
85
+ Test the apply_repetition_penalties custom op with num_seqs=0
86
+ against a reference implementation.
87
+ """
88
+ num_seqs = 0
89
+ vocab_size = 17
90
+ repetition_penalty = 1.05
91
+ dtype = torch.float32
92
+ seed = 0
93
+
94
+ current_platform.seed_everything(seed)
95
+ torch.set_default_device("cuda:0")
96
+
97
+ # Create test data
98
+ logits = torch.randn(num_seqs, vocab_size, dtype=dtype)
99
+
100
+ # Create masks with some random tokens marked as repeated
101
+ prompt_mask = torch.zeros(num_seqs, vocab_size, dtype=torch.bool)
102
+ output_mask = torch.zeros(num_seqs, vocab_size, dtype=torch.bool)
103
+
104
+ # No tokens to mark as repeated since num_seqs=0
105
+
106
+ # Create repetition penalties tensor
107
+ repetition_penalties = torch.full((num_seqs, ),
108
+ repetition_penalty,
109
+ dtype=dtype)
110
+
111
+ # Run all three implementations
112
+ logits_torch = logits.clone()
113
+ logits_cuda = logits.clone()
114
+
115
+ apply_repetition_penalties_torch(logits_torch, prompt_mask, output_mask,
116
+ repetition_penalties)
117
+ apply_repetition_penalties_cuda(logits_cuda, prompt_mask, output_mask,
118
+ repetition_penalties)
119
+
120
+ # Compare all outputs to reference
121
+ torch.testing.assert_close(logits_torch, logits_cuda, rtol=1e-3, atol=1e-3)
122
+
123
+ # Test the operator by applying the opcheck utility
124
+ opcheck(torch.ops._C.apply_repetition_penalties_,
125
+ (logits.clone(), prompt_mask, output_mask, repetition_penalties))