@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 130 lines 4.24 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5
6 from tests.kernels.utils import opcheck
7 from vllm._custom_ops import (
8 apply_repetition_penalties_cuda,
9 apply_repetition_penalties_torch,
10 )
11 from vllm.platforms import current_platform
12 from vllm.utils.torch_utils import set_random_seed
13
14 NUM_SEQS = [1, 2, 3, 4, 8, 13, 17, 32, 37, 256, 1023, 1024, 1025]
15 # [stress, stress, stress, Qwen, llama 4]
16 VOCAB_SIZES = [17, 256, 1019, 151936, 202048]
17 REPETITION_PENALTY_VALUES = [1.05]
18 SEEDS = [0]
19 DTYPES = [torch.float32, torch.float16]
20
21
22 @pytest.mark.parametrize("num_seqs", NUM_SEQS)
23 @pytest.mark.parametrize("vocab_size", VOCAB_SIZES)
24 @pytest.mark.parametrize("repetition_penalty", REPETITION_PENALTY_VALUES)
25 @pytest.mark.parametrize("dtype", DTYPES)
26 @pytest.mark.parametrize("seed", SEEDS)
27 @pytest.mark.skipif(
28 not current_platform.is_cuda(), reason="This test for checking CUDA kernel"
29 )
30 @torch.inference_mode()
31 def test_apply_repetition_penalties(
32 num_seqs: int,
33 vocab_size: int,
34 repetition_penalty: float,
35 dtype: torch.dtype,
36 seed: int,
37 ) -> None:
38 """Test the apply_repetition_penalties custom op
39 against a reference implementation.
40 """
41 set_random_seed(seed)
42 torch.set_default_device("cuda:0")
43
44 # Create test data
45 logits = torch.randn(num_seqs, vocab_size, dtype=dtype)
46
47 # Create masks with some random tokens marked as repeated
48 prompt_mask = torch.zeros(num_seqs, vocab_size, dtype=torch.bool)
49 output_mask = torch.zeros(num_seqs, vocab_size, dtype=torch.bool)
50
51 # Mark some tokens as repeated in prompt and output
52 prompt_indices = torch.randint(0, vocab_size, (num_seqs, max(1, vocab_size // 200)))
53 output_indices = torch.randint(0, vocab_size, (num_seqs, max(1, vocab_size // 200)))
54
55 for i in range(num_seqs):
56 prompt_mask[i, prompt_indices[i]] = True
57 output_mask[i, output_indices[i]] = True
58
59 # Create repetition penalties tensor
60 repetition_penalties = torch.full((num_seqs,), repetition_penalty, dtype=dtype)
61
62 # Run all three implementations
63 logits_torch = logits.clone()
64 logits_cuda = logits.clone()
65
66 apply_repetition_penalties_torch(
67 logits_torch, prompt_mask, output_mask, repetition_penalties
68 )
69 apply_repetition_penalties_cuda(
70 logits_cuda, prompt_mask, output_mask, repetition_penalties
71 )
72
73 # Compare all outputs to reference
74 torch.testing.assert_close(logits_torch, logits_cuda, rtol=1e-3, atol=1e-3)
75
76 # Test the operator by applying the opcheck utility
77 opcheck(
78 torch.ops._C.apply_repetition_penalties_,
79 (logits.clone(), prompt_mask, output_mask, repetition_penalties),
80 )
81
82
83 @pytest.mark.skipif(
84 not current_platform.is_cuda(), reason="This test for checking CUDA kernel"
85 )
86 @torch.inference_mode()
87 def test_apply_repetition_penalties_zero_seqs() -> None:
88 """Test the apply_repetition_penalties custom op with num_seqs=0
89 against a reference implementation.
90 """
91 num_seqs = 0
92 vocab_size = 17
93 repetition_penalty = 1.05
94 dtype = torch.float32
95 seed = 0
96
97 set_random_seed(seed)
98 torch.set_default_device("cuda:0")
99
100 # Create test data
101 logits = torch.randn(num_seqs, vocab_size, dtype=dtype)
102
103 # Create masks with some random tokens marked as repeated
104 prompt_mask = torch.zeros(num_seqs, vocab_size, dtype=torch.bool)
105 output_mask = torch.zeros(num_seqs, vocab_size, dtype=torch.bool)
106
107 # No tokens to mark as repeated since num_seqs=0
108
109 # Create repetition penalties tensor
110 repetition_penalties = torch.full((num_seqs,), repetition_penalty, dtype=dtype)
111
112 # Run all three implementations
113 logits_torch = logits.clone()
114 logits_cuda = logits.clone()
115
116 apply_repetition_penalties_torch(
117 logits_torch, prompt_mask, output_mask, repetition_penalties
118 )
119 apply_repetition_penalties_cuda(
120 logits_cuda, prompt_mask, output_mask, repetition_penalties
121 )
122
123 # Compare all outputs to reference
124 torch.testing.assert_close(logits_torch, logits_cuda, rtol=1e-3, atol=1e-3)
125
126 # Test the operator by applying the opcheck utility
127 opcheck(
128 torch.ops._C.apply_repetition_penalties_,
129 (logits.clone(), prompt_mask, output_mask, repetition_penalties),
130 )