@@ -138,6 +138,8 @@ def __init__(
138
138
self .text_creator = EndlessTextCreator (
139
139
data = config .source ,
140
140
)
141
+ # Add counter for unique prefixes
142
+ self .request_counter = 0
141
143
142
144
def __iter__ (
143
145
self ,
@@ -170,22 +172,46 @@ def __iter__(
170
172
output_tokens_sampler ,
171
173
):
172
174
start_index = rand .randint (0 , len (self .text_creator .words ))
175
+ # Increment counter for each request
176
+ self .request_counter += 1
173
177
yield {
174
- "prompt" : self ._create_prompt (prompt_tokens , start_index ),
178
+ "prompt" : self ._create_prompt (
179
+ prompt_tokens , start_index , self .request_counter
180
+ ),
175
181
"prompt_tokens_count" : prompt_tokens ,
176
182
"output_tokens_count" : output_tokens ,
177
183
}
178
184
179
- def _create_prompt (self , prompt_tokens : int , start_index : int ) -> str :
185
+ def _create_prompt (
186
+ self , prompt_tokens : int , start_index : int , request_id : int
187
+ ) -> str :
188
+ """
189
+ Create a prompt with unique prefix to prevent vLLM prefix caching.
190
+ Args:
191
+ prompt_tokens: Target number of tokens for the prompt
192
+ start_index: Starting position in the text corpus
193
+ request_id: Unique identifier for this request (used as prefix)
194
+ Returns:
195
+ Generated prompt string with unique prefix
196
+ """
180
197
if prompt_tokens <= 0 :
181
- return ""
198
+ return f"{ request_id } : "
199
+
200
+ unique_prefix = f"{ request_id } : "
201
+
202
+ # Calculate how many tokens the prefix uses
203
+ prefix_tokens = len (self .processor .tokenize (unique_prefix ))
204
+
205
+ # Adjust target tokens to account for the prefix
206
+ remaining_tokens = max (1 , prompt_tokens - prefix_tokens )
182
207
183
208
left = start_index
184
- right = start_index + 4 * prompt_tokens
209
+ right = start_index + 4 * remaining_tokens
185
210
186
211
while left < right :
187
212
mid = (left + right ) // 2
188
- test_prompt = self .text_creator .create_text (start_index , mid - start_index )
213
+ base_text = self .text_creator .create_text (start_index , mid - start_index )
214
+ test_prompt = unique_prefix + base_text
189
215
test_tokens = len (self .processor .tokenize (test_prompt ))
190
216
191
217
if test_tokens == prompt_tokens :
@@ -195,7 +221,8 @@ def _create_prompt(self, prompt_tokens: int, start_index: int) -> str:
195
221
else :
196
222
right = mid
197
223
198
- return self .text_creator .create_text (start_index , left - start_index )
224
+ base_text = self .text_creator .create_text (start_index , left - start_index )
225
+ return unique_prefix + base_text
199
226
200
227
201
228
class SyntheticDatasetCreator (DatasetCreator ):
0 commit comments