@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7e6544c797

Fork 0

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 14, 05:30 UTC

[Perf] Parallelize fill_bitmask to accelerate high-throughput guided decoding (#21862)

Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai>

Benjamin Chislett committed Aug 5, 2025 at 22:57 UTC 7e6544c7978364fcb8178f4ab8b1325e45880aa9

3 files changed +105 -42

vllm/v1/structured_output/__init__.py

+92 -39

@@ -3,7 +3,7 @@

3 from __future__ import annotations

5 import multiprocessing

6 -from concurrent.futures import ThreadPoolExecutor

6 +from concurrent.futures import Future, ThreadPoolExecutor

7 from typing import TYPE_CHECKING, Optional

9 from vllm.config import VllmConfig

@@ -40,6 +40,17 @@ class StructuredOutputManager:

40 self._grammar_bitmask: Optional[torch.Tensor] = None

41 self._full_mask = torch.tensor(-1, dtype=torch.int32)

43 + max_batch_size = self.vllm_config.scheduler_config.max_num_seqs

44 + self.fill_bitmask_parallel_threshold = 128

45 + if self.fill_bitmask_parallel_threshold < max_batch_size:

46 + self.fill_bitmask_parallel_batch_size = 16

47 + # Use:

48 + # - at least 1 CPU

49 + # - at most half the number of CPUs or 8, whichever is less

50 + max_workers = max(1, min(multiprocessing.cpu_count() // 2, 8))

51 + self.executor_for_fillmask = ThreadPoolExecutor(

52 + max_workers=max_workers)

53 +

54 if not self.vllm_config.model_config.skip_tokenizer_init:

55 # The default max_workers if not specified is the number of

56 # CPUs * 5, which is way too high since these tasks are CPU-bound,

@@ -120,6 +131,26 @@ class StructuredOutputManager:

131 assert self.backend is not None

132 return self.backend.compile_grammar(request_type, grammar_spec)

133

134 + def _fill_bitmasks(

135 + self,

136 + batch: list[tuple[StructuredOutputGrammar, int, bool]],

137 + ) -> None:

138 + assert self._grammar_bitmask is not None

139 + for grammar, index, apply_bitmask in batch:

140 + if apply_bitmask and not grammar.is_terminated():

141 + grammar.fill_bitmask(self._grammar_bitmask, index)

142 + else:

143 + # Note that for thinking support, we will need to

144 + # reset the relevant part of the bitmask for consequent

145 + # requests here.

146 + self._grammar_bitmask[index].fill_(self._full_mask)

147 +

148 + def _async_submit_fill_bitmask(

149 + self,

150 + batch: list[tuple[StructuredOutputGrammar, int, bool]],

151 + ) -> Future:

152 + return self.executor_for_fillmask.submit(self._fill_bitmasks, batch)

153 +

154 def grammar_bitmask(

155 self,

156 requests: dict[str, Request],

@@ -146,7 +177,6 @@ class StructuredOutputManager:

177 self.backend.allocate_token_bitmask(

178 max_batch_size * (1 + max_num_spec_tokens))

179

149 - bitmask_tensor = self._grammar_bitmask

180 # Generate a batched bitmask for all structured output requests.

181 # When speculative decoding is enabled, we need to include multiple

182 # masks for each request, one for each possible bonus token position.

@@ -155,47 +185,61 @@ class StructuredOutputManager:

185 ordered_seq = sorted(structured_output_request_ids.items(),

186 key=lambda x: x[1])

187

158 - # Note that for thinking support, we will need to

159 - # reset the relevant part of the bitmask for consequent

160 - # request here.

161 - bitmask_tensor[:(len(ordered_seq) * (1 + max_num_spec_tokens))].fill_(

162 - self._full_mask)

163 -

164 - # NOTE: This outer loop can likely be parallelized to improve

165 - # performance of bitmask generation for large batches.

166 - for req_id, _ in ordered_seq:

167 - request = requests[req_id]

168 - structured_output_request = request.structured_output_request

169 -

170 - if TYPE_CHECKING:

171 - assert structured_output_request is not None

172 - assert structured_output_request.grammar is not None

173 - apply_bitmask: bool = True

174 - if self.reasoner is not None:

175 - if structured_output_request.reasoning_ended is None:

176 - structured_output_request.reasoning_ended = \

177 - self.reasoner.is_reasoning_end(request.prompt_token_ids)

178 - apply_bitmask = structured_output_request.reasoning_ended

179 -

180 - state_advancements = 0

181 - req_tokens = scheduled_spec_decode_tokens.get(req_id, []) + [None]

182 - for i, token in enumerate(req_tokens):

183 - if apply_bitmask and not \

184 - structured_output_request.grammar.is_terminated():

185 - structured_output_request.grammar.fill_bitmask(

186 - bitmask_tensor, cumulative_index)

187 - if token is not None:

188 - # In order to generate the correct bitmask for each

189 - # position in the speculative sequence, we advance

190 - # the FSM state for each speculative token and rollback

191 - # to restore the previous state when we are finished.

188 + # Optimized parallel filling of bitmasks for

189 + # non-spec, large-batch-size cases

190 + if len(ordered_seq) > self.fill_bitmask_parallel_threshold and \

191 + max_num_spec_tokens == 0:

192 + promises = []

193 + batch = []

194 + for req_id, _ in ordered_seq:

195 + request = requests[req_id]

196 + structured_output_request = request.structured_output_request

197 + if TYPE_CHECKING:

198 + assert structured_output_request is not None

199 + assert structured_output_request.grammar is not None

200 +

201 + apply_bitmask = self.should_fill_bitmask(request)

202 + batch.append((structured_output_request.grammar,

203 + cumulative_index, apply_bitmask))

204 + if len(batch) == self.fill_bitmask_parallel_batch_size:

205 + promises.append(self._async_submit_fill_bitmask(batch))

206 + batch = []

207 +

208 + cumulative_index += 1

209 + if batch:

210 + promises.append(self._async_submit_fill_bitmask(batch))

211 +

212 + # Wait for all bitmask filling tasks to complete.

213 + for promise in promises:

214 + promise.result()

215 + else:

216 + # Fallback to serial filling of bitmasks for small-batch-size cases

217 + for req_id, _ in ordered_seq:

218 + request = requests[req_id]

219 + structured_output_request = request.structured_output_request

220 +

221 + if TYPE_CHECKING:

222 + assert structured_output_request is not None

223 + assert structured_output_request.grammar is not None

224 + apply_bitmask = self.should_fill_bitmask(request)

225 +

226 + state_advancements = 0

227 + req_tokens = scheduled_spec_decode_tokens.get(req_id, [])

228 + for i, token in enumerate(req_tokens + [None]):

229 + self._fill_bitmasks([(structured_output_request.grammar,

230 + cumulative_index, apply_bitmask)])

231 +

232 + if apply_bitmask and token is not None and \

233 + not structured_output_request.grammar.is_terminated():

234 assert structured_output_request.grammar.accept_tokens(

235 req_id, [token])

236 state_advancements += 1

195 - cumulative_index += 1

196 - if state_advancements > 0:

197 - structured_output_request.grammar.rollback(state_advancements)

237 + cumulative_index += 1

238 + if state_advancements > 0:

239 + structured_output_request.grammar.rollback(

240 + state_advancements)

241

242 + bitmask_tensor = self._grammar_bitmask

243 if cumulative_index < bitmask_tensor.shape[0]:

244 bitmask_tensor = bitmask_tensor[:cumulative_index]

245

@@ -204,6 +248,15 @@ class StructuredOutputManager:

248 # and deserialization when sending this to the GPU workers.

249 return bitmask_tensor.numpy()

250

251 + def should_fill_bitmask(self, request: Request) -> bool:

252 + if self.reasoner is not None:

253 + assert request.structured_output_request is not None

254 + if request.structured_output_request.reasoning_ended is None:

255 + request.structured_output_request.reasoning_ended = \

256 + self.reasoner.is_reasoning_end(request.prompt_token_ids)

257 + return request.structured_output_request.reasoning_ended

258 + return True

259 +

260 def should_advance(self, request: Request) -> bool:

261 if not request.use_structured_output:

262 return False

vllm/v1/structured_output/backend_xgrammar.py

+6 -1

@@ -148,6 +148,7 @@ class XgrammarGrammar(StructuredOutputGrammar):

148 repr=False,

149 hash=False,

150 init=False)

151 + _is_terminated: bool = field(default=False, repr=False, hash=False)

152

153 def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:

154 """Accepts a list of tokens and advances the FSM.

@@ -155,6 +156,8 @@ class XgrammarGrammar(StructuredOutputGrammar):

156 Returns True if the FSM was advanced successfully.

157 Returns False if the FSM failed to advance.

158 """

159 + if self._is_terminated:

160 + return False

161 for token in tokens:

162 if not self.matcher.accept_token(token):

163 logger.error(

@@ -162,6 +165,7 @@ class XgrammarGrammar(StructuredOutputGrammar):

165 "for tokens %s. Please file an issue.", request_id, token)

166 return False

167 self.num_processed_tokens += 1

168 + self._is_terminated = self.matcher.is_terminated()

169 return True

170

171 def validate_tokens(self, tokens: list[int]) -> list[int]:

@@ -184,12 +188,13 @@ class XgrammarGrammar(StructuredOutputGrammar):

188 def rollback(self, num_tokens: int) -> None:

189 self.matcher.rollback(num_tokens)

190 self.num_processed_tokens -= num_tokens

191 + self._is_terminated = self.matcher.is_terminated()

192

193 def fill_bitmask(self, bitmask: torch.Tensor, idx: int) -> None:

194 self.matcher.fill_next_token_bitmask(bitmask, idx)

195

196 def is_terminated(self) -> bool:

192 - return self.matcher.is_terminated()

197 + return self._is_terminated

198

199 def reset(self):

200 self.num_processed_tokens = 0

vllm/v1/worker/gpu_model_runner.py

+7 -2

@@ -1324,9 +1324,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1324 cumulative_index += 1 + num_spec_tokens

1325 grammar_bitmask = sorted_bitmask

1326

1327 + # If the grammar bitmask and the logits have the same shape

1328 + # we don't need to pass indices to the kernel,

1329 + # since the bitmask is already aligned with the logits.

1330 + skip_out_indices = grammar_bitmask.shape[0] == logits.shape[0]

1331 +

1332 # Serialization of np.ndarray is much more efficient than a tensor,

1333 # so we receive it in that format.

1329 - grammar_bitmask = torch.from_numpy(grammar_bitmask)

1334 + grammar_bitmask = torch.from_numpy(grammar_bitmask).contiguous()

1335

1336 # Force use of the torch.compile implementation from xgrammar to work

1337 # around issues with the Triton kernel in concurrent structured output

@@ -1334,7 +1339,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1339 xgr_torch_compile.apply_token_bitmask_inplace_torch_compile(

1340 logits,

1341 grammar_bitmask.to(self.device, non_blocking=True),

1337 - indices=out_indices,

1342 + indices=out_indices if not skip_out_indices else None,

1343 )

1344

1345 def sync_and_slice_intermediate_tensors(